text:
# ensure in right directory and environment
cd(dirname(@__FILE__))
using Pkg
pkg"activate ."

##

using CMBLensing
using Documenter

# convert Jupyter notebooks to markdown and do some simple preprocessing before
# feeding into Documenter.jl

function convert_to_markdown(file)
    run(`jupyter nbconvert src/$file --to markdown --template documenter.tpl --output-dir src-staging`)
    return "src-staging/$(replace(file, "ipynb"=>"md"))"
end

function convert_equations(file)
    contents = read(file, String)
    contents = replace(contents, r"\$\$(.*?)\$\$"s => s"""```math
    \g<1>
    ```""")
    contents = replace(contents, r"\* \$(.*?)\$" => s"* ``\g<1>``") # starting a line with inline math screws up tex2jax for some reason
    write(file, contents)
    return file
end

for file in readdir("src")
    if endswith(file, "ipynb")
        file |> convert_to_markdown |> convert_equations
    elseif !startswith(file, ".")
        cp("src/$file", "src-staging/$file", force=true)
    end
end
rm("src-staging/index.md",force=true)
symlink("../../README.md","src-staging/index.md")

# highlight output cells (i.e. anything withouout a language specified) white

@eval Documenter.Writers.HTMLWriter function mdconvert(c::Markdown.Code, parent::MDBlockContext; kwargs...)
    @tags pre code
    language = isempty(c.language) ? "none" : c.language
    pre[".language-$(language)"](code[".language-$(language)"](c.code))
end

# adds the MyBinder button on each page that is a notebook

@eval Documenter.Writers.HTMLWriter function render_navbar(ctx, navnode, edit_page_link::Bool)
    @tags div header nav ul li a span img

    # The breadcrumb (navigation links on top)
    navpath = @show Documents.navpath(navnode)
    header_links = map(navpath) do nn
        title = mdconvert(pagetitle(ctx, nn); droplinks=true)
        nn.page === nothing ? li(a[".is-disabled"](title)) : li(a[:href => navhref(ctx, nn, navnode)](title))
    end
    header_links[end] = header_links[end][".is-active"]
    breadcrumb = nav[".breadcrumb"]() ### modified to not show breadcrumbs

    # The "Edit on GitHub" links and the hamburger to open the sidebar (on mobile) float right
    navbar_right = div[".docs-right"]

    ### custom code to add MyBinder link
    if edit_page_link
        pageurl = get(getpage(ctx, navnode).globals.meta, :EditURL, getpage(ctx, navnode).source)
        nbpath = foldl(replace,["src-staging"=>"src",".md"=>".ipynb"], init=pageurl)
        if @show isfile(nbpath)
            url = "https://mybinder.org/v2/gh/marius311/CMBLensing.jl/gh-pages?urlpath=lab/tree/$(basename(nbpath))"
            push!(navbar_right.nodes, a[".docs-right", :href => url](img[:src => "https://mybinder.org/badge_logo.svg"]()))
        end
    end
    ###

    # Set the logo and name for the "Edit on.." button.
    if edit_page_link && (ctx.settings.edit_link !== nothing) && !ctx.settings.disable_git
        host_type = Utilities.repo_host_from_url(ctx.doc.user.repo)
        if host_type == Utilities.RepoGitlab
            host = "GitLab"
            logo = "\uf296"
        elseif host_type == Utilities.RepoGithub
            host = "GitHub"
            logo = "\uf09b"
        elseif host_type == Utilities.RepoBitbucket
            host = "BitBucket"
            logo = "\uf171"
        else
            host = ""
            logo = "\uf15c"
        end
        hoststring = isempty(host) ? " source" : " on $(host)"

        pageurl = get(getpage(ctx, navnode).globals.meta, :EditURL, getpage(ctx, navnode).source)

        edit_branch = isa(ctx.settings.edit_link, String) ? ctx.settings.edit_link : nothing
        url = if Utilities.isabsurl(pageurl)
            pageurl
        else
            if !(pageurl == getpage(ctx, navnode).source)
                # need to set users path relative the page itself
                pageurl = joinpath(first(splitdir(getpage(ctx, navnode).source)), pageurl)
            end
            Utilities.url(ctx.doc.user.repo, pageurl, commit=edit_branch)
        end
        if url !== nothing
            edit_verb = (edit_branch === nothing) ? "View" : "Edit"
            title = "$(edit_verb)$hoststring"
            push!(navbar_right.nodes,
                a[".docs-edit-link", :href => url, :title => title](
                    span[".docs-icon.fab"](logo),
                    span[".docs-label.is-hidden-touch"](title)
                )
            )
        end
    end

    # Settings cog
    push!(navbar_right.nodes, a[
        "#documenter-settings-button.docs-settings-button.fas.fa-cog",
        :href => "#", :title => "Settings",
    ])

    # Hamburger on mobile
    push!(navbar_right.nodes, a[
        "#documenter-sidebar-button.docs-sidebar-button.fa.fa-bars.is-hidden-desktop",
        :href => "#"
    ])

    # Construct the main <header> node that should be the first element in div.docs-main
    header[".docs-navbar"](breadcrumb, navbar_right)
end

makedocs(
    sitename="CMBLensing.jl",
    source = "src-staging",
    format = Documenter.HTML(
        assets = ["assets/cmblensing.css"],
        disable_git = true,
    ),
    pages = [
        "index.md",
        "01_lense_a_map.md",
        "02_posterior.md",
        "03_joint_MAP_example.md",
        "04_from_python.md",
        "05_field_basics.md",
        "06_gpu.md",
        "api.md"
    ],
)

if haskey(ENV, "IMAGE_NAME")
    open("build/Dockerfile","w") do io
        write(io,"FROM $(ENV["IMAGE_NAME"])")
    end
end

deploydocs(
    repo = "github.com/marius311/CMBLensing.jl.git",
    push_preview = true,
    forcepush = true
)


text:
# function to determine whether an x, y point is in the unit circle
function in_circle(x_pos::Float64, y_pos::Float64, radius::Float64)
    if (x_pos^2 + y_pos^2 < radius^2)
        return true
    else
        return false
    end
end

# function to integrate a unit circle to find pi via monte_carlo
function monte_carlo(n::Int64, radius::Float64)

    pi_count = 0
    for i = 1:n
        point_x = rand()
        point_y = rand()

        if (in_circle(point_x, point_y, radius))
            pi_count += 1
        end
    end

    pi_estimate = 4*pi_count/(n*radius^2)
    println("Percent error is: ", signif(100*(pi - pi_estimate)/pi, 3), " %")
end

monte_carlo(10000000, 0.5)

ext: jl


text:
function WriteImage(path, dm, du)
    (m, n) = size(dm)
    fid = open(path, "w")
    write(fid, Int32(m), Int32(n))
    write(fid, convert(Array{Float32},vec(dm)))
    write(fid, convert(Array{Float32},vec(du)))
    close(fid)
end

function ReadImage(path)
    fid = open(path, "r")
    m = Int64(read(fid, Int32))
    n = Int64(read(fid, Int32))
    dm = read(fid, Float32, m*n)
    dm = convert(Array{Float64}, dm)
    du = read(fid, Float32, m*n)
    du = convert(Array{Float64}, du)
    dm = reshape(dm, m, n)
    du = reshape(du, m, n)
    close(fid)
    return dm, du
end

function laplaceFilter(I)
    (m, n) = size(I)
    d = zeros(typeof(I[1]), m, n)
    for ix = 2 : n-1
        for iz = 2 : m-1
            d[iz, ix] = I[iz-1,ix]+I[iz+1,ix]+I[iz,ix-1]+I[iz,ix+1]-4*I[iz,ix]
        end
    end
    return d
end

function lambdamu2vpvs(dm, du, vps, vss, rho)
    (m, n) = size(dm)
    dvp = zeros(typeof(dm[1]), m, n)
    dvp = 2* vps .* rho .* dm
    dvs = -4 * vss .* rho .* dm + 2 * vss .* rho .* du
    return dvp, dvs
end

function envelope(d)
    (nt, nx) = size(d)
    dout     = zeros(nt, nx)
    for ix =1 : nx
        temp = hilbert(d[:, ix])
        dout[:, ix] = abs(temp)
    end
    return dout
end

function gec(din::Array{Float64,2})
    (nt, nx) = size(din)
    dout     = zeros(nt, nx)
    nop      = 2*int(nt/10) + 1
    smoother = triang(nop)
    trip     = int(nop / 20)
    env = envelope(din)
    for ix =1 : nx
        temp = env[:, ix]
        temp = conv(temp, smoother)
        envsm= temp[int(round(nop/2)) : nt+int(round(nop)/2)-1]
        envsm[1:trip] = envsm[trip+1]
        envsm[nt-trip+1: end] = envsm[nt-trip]
        trin  = din[:, ix]
        trout = trin ./ envsm
        trout = trout / maximum(abs(trout))
        dout[:, ix] = trout
    end
    return dout
end

max_stars_repo_name: WneleiGao/ElasticWave


text:
function simulate!(rng::AbstractRNG,
                   sim::Simulation;
                   tmax::Float64=Inf,
                   nmax::Int64=100000,
                   pmax::Float64=Inf)
  if tmax <= 0.0
    @error "The simulation time maximum must be > 0.0"
  elseif nmax < 1
    @error "The iteration maximum must be > 0"
  elseif pmax <= 0.0
    @error "The processing time maximum must be > 0.0"
  else
    @debug "Simulating until: \n* a simulation time maximum of $tmax,\n* an iteration maximum of $nmax, or,\n* a processing time maximum of $pmax seconds.\n"
  end
  stoptime = time() + pmax
  while true
    # Generate the new event and associated specific transmission when applicable
    event = generate(rng, Event, sim.event_rates, sim.time)
    transmission = generate(rng, Transmission, sim.transmission_rates, event)
    if _time(event) > tmax
      sim.time = tmax
      @debug "Simulation stopped: simulation time maximum reached"
      break
    else
      update!(sim, event, transmission)
    end
    if sim.iterations >= nmax
      @debug "Simulation stopped: iteration maximum reached"
      break
    elseif pmax < Inf && time() >= stoptime
      @debug "Simulation stopped: processing time maximum reached"
      break
    end
  end
  return sim
end

simulate!(sim::Simulation; kwargs...) = simulate!(default_rng(), sim; kwargs...)


text:
__precompile__()

"""
    module BEDMatrices1.0

A collection of tools for efficiently reading or indexing into .bed
files for calculation or other manipulations. Functionality primarily
revolves around the `immutable` `BEDMatrix`; see `BEDMatrix`
documentation for more details.

"""
module BEDMatrices1.0
using Pkg
using LinearAlgebra

module Consts
using Pkg
using LinearAlgebra
include("constants.jl")
end

const NA_byte = Consts.NA_byte

include("bedmatrix.jl")
include("coltools.jl")
include("gwas.jl")

export BEDintomatrix,
    BEDintomatrix!,
    BEDMatrix,
    path,
    rownames,
    colnames,
    getflips,
    setflips!,
    NArep,
    getquartermap,
    hasNAs,
    countNAs,
    BEDdot,
    column_dist,
    column_sum,
    column_dot,
    column_NAsup_dot,
    column_sumabs2,
    column_norm,
    column_mean_and_std,
    column_mean_and_var,
    GWAS,
    gwas_writecsv

end


text:
export AbstractEnergyModel, TotalEnergyModel, θModel
#### Energy component in atmosphere model
abstract type AbstractEnergyModel end
struct TotalEnergyModel <: AbstractEnergyModel end
struct θModel <: AbstractEnergyModel end

vars_state(::TotalEnergyModel, ::AbstractStateType, FT) = @vars()
vars_state(::TotalEnergyModel, ::Prognostic, FT) = @vars(ρe::FT)
vars_state(::TotalEnergyModel, ::Gradient, FT) = @vars(h_tot::FT)
vars_state(::TotalEnergyModel, ::GradientFlux, FT) =
    @vars(∇h_tot::SVector{3, FT})
vars_state(::θModel, ::AbstractStateType, FT) = @vars()
vars_state(::θModel, ::Prognostic, FT) = @vars(ρθ_liq_ice::FT)
vars_state(::θModel, ::Gradient, FT) = @vars(θ_liq_ice::FT)
vars_state(::θModel, ::GradientFlux, FT) = @vars(∇θ_liq_ice::SVector{3, FT})

vars_state_filtered(::θModel, FT) = @vars(θ_liq_ice::FT)
vars_state_filtered(::TotalEnergyModel, FT) = @vars(e::FT)

function compute_gradient_argument!(
    energy::TotalEnergyModel,
    atmos::AtmosModel,
    transform,
    state,
    aux,
    t,
)
    ts = recover_thermo_state(atmos, state, aux)
    e_tot = state.energy.ρe * (1 / state.ρ)
    transform.energy.h_tot = total_specific_enthalpy(ts, e_tot)
end

function compute_gradient_argument!(
    energy::θModel,
    ::AtmosModel,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform.energy.θ_liq_ice = state.energy.ρθ_liq_ice / state.ρ
end

compute_gradient_flux!(::TotalEnergyModel, _...) = nothing

function compute_gradient_flux!(::θModel, diffusive, ∇transform, state, aux, t)
    diffusive.energy.∇θ_liq_ice = ∇transform.energy.θ_liq_ice
end

function compute_gradient_flux!(
    ::TotalEnergyModel,
    diffusive,
    ∇transform,
    state,
    aux,
    t,
)
    diffusive.energy.∇h_tot = ∇transform.energy.h_tot
end


text:
using NCDatasets
using Formatting
using SharedArrays

getFn_output = (varname, ens) -> format("xtt_b.e11.B1850C5CN.f09_g16.005.pop.h.{}.{:02d}0001-{:02d}9912.nc", varname, ens, ens)
getFn_input  = (varname, ens) -> format("b.e11.B1850C5CN.f09_g16.005.pop.h.{}.{:02d}0001-{:02d}9912.nc", varname, ens, ens)
genAttr = function(
    units,
    long_name;
    extra :: Dict = Dict()
)
    d = Dict(
        "units" => units,
        "long_name" => long_name,
    )

    for (k, v) in extra
        d[k] = v
    end

    return d
end

attr2Dict = function(attr)
    d = Dict()
    for k in keys(attr)
        d[k] = attr[k]
    end

    return d
end

output_file = getFn_output("AVGSALT_AVGTEMP", 5)

h_ML = 30.0
ds_dom = Dataset("domain.ocn.gx1v6.090206.nc", "r")
Nx = ds_dom.dim["ni"]
Ny = ds_dom.dim["nj"]
mask = ds_dom["mask"][:]

qflx = SharedArray{Float64}(Nx, Ny, 12)
qflx .= 0.0

ens_list = [7, 8, 9, 10]
year_cnt = 0
for ens in ens_list

    println("Doing Ensemble:", ens)

    ds_SHF = Dataset(getFn("SHF", ens), "r")
    ds_SST = Dataset(getFn("SST", ens), "r")

    F   = nomissing(ds_SHF["SHF"][:], NaN)
    SST = view(nomissing(ds_SST["SST"][:], NaN), :, :, 1, :)

    Nt = ds_SHF.dim["time"]

    if Nt % 12 != 0
        throw(ErrorException(format("Time is not multiple of 12. The length I got is {}.", Nt)))
    end

    println("Calculating...")
    @time @sync @distributed for idx in CartesianIndices((1:Nx, 1:Ny))
        i = idx[1]
        j = idx[2]
        for t = 13:Nt-12
            if mask[i, j] == 0
                qflx[i, j, :] .= NaN
                continue
            end

            mon = mod(t-1, 12) + 1

            beg_t = t-1
            end_t = t

            qflx[i, j, mon] += h_ML * ρ * c_p * (SST[i, j, end_t] - SST[i, j, beg_t]) / Δts[mon] - (F[i, j, beg_t] + F[i, j, end_t]) / 2.0
        end
    end
    println("done.")
    close(ds_SHF)
    close(ds_SST)

    global year_cnt += (Nt / 12.0)

end

qflx ./= year_cnt

missing_value = 1e20
qflx[isnan.(qflx)] .= missing_value

Dataset(output_file, "c") do ds

    defDim(ds, "Nx", Nx)
    defDim(ds, "Ny", Ny)
    defDim(ds, "time", 12)

    defVar(ds, "time", ts, ("time",); attrib=Dict("calendar"=>"noleap", "units"=>"days since 0001-01-01 00:00:00"))

    for varname in ["mask", "area", "xc", "yc", "frac"]
        vardata = ds_dom[varname]
        dims = ("Nx", "Ny")
        v = nomissing(vardata[:], -1)
        v = defVar(ds, varname, v , dims; attrib=attr2Dict(vardata.attrib))
    end

    v = defVar(ds, "qdp", qflx, ("Nx", "Ny", "time");
        fillvalue=missing_value,
        attrib=Dict(
            "long_name" => "Temperature flux convergence in energy",
            "units"     => "W / m^2"
        )
    )
end

close(ds_dom)

max_stars_repo_name: meteorologytoday/CESM-diagnostic

max_issues_repo_name: meteorologytoday/IOM

max_forks_repo_name: meteorologytoday/IOM


text:
@testset "lz4frame" begin
     testIn = "Far out in the uncharted backwaters of the unfashionable end of the west-
 ern  spiral  arm  of  the  Galaxy  lies  a  small  unregarded  yellow  sun."
    test_size = convert(UInt, length(testIn))
    version = CodecLz4.LZ4F_getVersion()

    @testset "Errors" begin
        no_error = UInt(0)
        @test !CodecLz4.LZ4F_isError(no_error)
        @test CodecLz4.LZ4F_getErrorName(no_error) == "Unspecified error code"

        ERROR_GENERIC = typemax(UInt)
        @test CodecLz4.LZ4F_isError(ERROR_GENERIC)
        @test CodecLz4.LZ4F_getErrorName(ERROR_GENERIC) == "ERROR_GENERIC"
    end

    @testset "keywords" begin

        frame = CodecLz4.LZ4F_frameInfo_t()
        @test frame.blockSizeID == Cuint(default_size)
        @test frame.blockMode == Cuint(block_linked)
        @test frame.contentChecksumFlag == Cuint(0)
        @test frame.frameType == Cuint(normal_frame)
        @test frame.contentSize == Culonglong(0)
        @test frame.dictID == Cuint(0)
        @test frame.blockChecksumFlag == Cuint(0)

        prefs = CodecLz4.LZ4F_preferences_t(frame)

        @test prefs.frameInfo == frame
        @test prefs.compressionLevel == Cint(0)
        @test prefs.autoFlush == Cuint(0)
        @test prefs.reserved == (Cuint(0), Cuint(0), Cuint(0), Cuint(0))

        frame = CodecLz4.LZ4F_frameInfo_t(
            blocksizeid = max64KB,
            blockmode = block_independent,
            contentchecksum = true,
            blockchecksum = true,
            frametype = skippable_frame,
            contentsize = 100
            )

        @test frame.blockSizeID == Cuint(4)
        @test frame.blockMode == Cuint(1)
        @test frame.contentChecksumFlag == Cuint(1)
        @test frame.frameType == Cuint(1)
        @test frame.contentSize == Culonglong(100)
        @test frame.blockChecksumFlag == Cuint(1)

        prefs = CodecLz4.LZ4F_preferences_t(frame, compressionlevel=5, autoflush = true)

        @test prefs.frameInfo == frame
        @test prefs.compressionLevel == Cint(5)
        @test prefs.autoFlush == Cuint(1)
        @test prefs.reserved == (Cuint(0), Cuint(0), Cuint(0), Cuint(0))

    end

    @testset "CompressionCtx" begin
        ctx = Ref{Ptr{CodecLz4.LZ4F_cctx}}(C_NULL)

        @test_nowarn err = CodecLz4.LZ4F_createCompressionContext(ctx, version)
        @test err == 0

        @test_nowarn CodecLz4.check_context_initialized(ctx[])

        err = CodecLz4.LZ4F_freeCompressionContext(ctx[])
        @test err == 0
        @test !CodecLz4.LZ4F_isError(err)

        ctx = Ptr{CodecLz4.LZ4F_cctx}(C_NULL)
        @test_throws CodecLz4.LZ4Exception CodecLz4.check_context_initialized(ctx)
    end

    @testset "DecompressionCtx" begin
        dctx = Ref{Ptr{CodecLz4.LZ4F_dctx}}(C_NULL)

        @test_nowarn err = CodecLz4.LZ4F_createDecompressionContext(dctx, version)
        @test err == 0

        @test_nowarn CodecLz4.check_context_initialized(dctx[])

        @test_nowarn CodecLz4.LZ4F_resetDecompressionContext(dctx[])

        err = CodecLz4.LZ4F_freeDecompressionContext(dctx[])
        @test err == 0

        dctx = Ptr{CodecLz4.LZ4F_dctx}(C_NULL)
        @test_throws CodecLz4.LZ4Exception CodecLz4.check_context_initialized(dctx)
        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_resetDecompressionContext(dctx)

    end

    function test_decompress(origsize, buffer)
        @testset "Decompress" begin
            dctx = Ref{Ptr{CodecLz4.LZ4F_dctx}}(C_NULL)
            srcsize = Ref{Csize_t}(origsize)
            dstsize =  Ref{Csize_t}(8*1280)
            decbuffer = Vector{UInt8}(undef, 1280)

            frameinfo = Ref(CodecLz4.LZ4F_frameInfo_t())

            @test_nowarn err = CodecLz4.LZ4F_createDecompressionContext(dctx, version)

            @test_nowarn result = CodecLz4.LZ4F_getFrameInfo(dctx[], frameinfo, buffer, srcsize)
            @test srcsize[] > 0

            offset = srcsize[]
            srcsize[] = origsize - offset

            @test_nowarn result = CodecLz4.LZ4F_decompress(dctx[], decbuffer, dstsize, pointer(buffer) + offset, srcsize, C_NULL)
            @test srcsize[] > 0

            @test testIn == unsafe_string(pointer(decbuffer), dstsize[])

            result = CodecLz4.LZ4F_freeDecompressionContext(dctx[])
            @test !CodecLz4.LZ4F_isError(result)
        end

    end

    function test_invalid_decompress(origsize, buffer)
        @testset "DecompressInvalid" begin

            dctx = Ref{Ptr{CodecLz4.LZ4F_dctx}}(C_NULL)
            srcsize = Ref{Csize_t}(origsize)
            dstsize =  Ref{Csize_t}(1280)
            decbuffer = Vector{UInt8}(undef, 1280)

            frameinfo = Ref(CodecLz4.LZ4F_frameInfo_t())

            CodecLz4.LZ4F_createDecompressionContext(dctx, version)

            buffer[1:CodecLz4.LZ4F_HEADER_SIZE_MAX] .= 0x10
            @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_getFrameInfo(dctx[], frameinfo, buffer, srcsize)

            offset = srcsize[]
            srcsize[] = origsize - offset

            @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_decompress(dctx[], decbuffer, dstsize, pointer(buffer) + offset, srcsize, C_NULL)

            result = CodecLz4.LZ4F_freeDecompressionContext(dctx[])
            @test !CodecLz4.LZ4F_isError(result)
        end
    end

    @testset "Compress" begin
        ctx = Ref{Ptr{CodecLz4.LZ4F_cctx}}(C_NULL)
        err = CodecLz4.LZ4F_isError(CodecLz4.LZ4F_createCompressionContext(ctx, version))
        @test !err

        prefs = Ptr{CodecLz4.LZ4F_preferences_t}(C_NULL)

        bound = CodecLz4.LZ4F_compressBound(test_size, prefs)
        @test bound > 0

        bufsize = bound + CodecLz4.LZ4F_HEADER_SIZE_MAX
        buffer = Vector{UInt8}(undef, ceil(Int, bound / 8))

        @test_nowarn result = CodecLz4.LZ4F_compressBegin(ctx[], buffer, bufsize, prefs)

        offset = result
        @test_nowarn result = CodecLz4.LZ4F_compressUpdate(ctx[], pointer(buffer) + offset, bufsize - offset, pointer(testIn), test_size, C_NULL)

        offset += result
        @test_nowarn result = CodecLz4.LZ4F_flush(ctx[], pointer(buffer)+offset, bufsize - offset, C_NULL)

        offset += result
        @test_nowarn result = CodecLz4.LZ4F_compressEnd(ctx[], pointer(buffer)+offset, bufsize - offset, C_NULL)
        @test result > 0

        offset += result

        result = CodecLz4.LZ4F_freeCompressionContext(ctx[])
        @test !CodecLz4.LZ4F_isError(result)

        test_decompress(offset, buffer)
        test_invalid_decompress(offset, buffer)
    end

    @testset "CompressUninitialized" begin
        ctx = Ref{Ptr{CodecLz4.LZ4F_cctx}}(C_NULL)

        prefs = Ptr{CodecLz4.LZ4F_preferences_t}(C_NULL)

        bufsize = test_size
        buffer = Vector{UInt8}(undef, test_size)

        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_compressBegin(ctx[], buffer, bufsize, prefs)
        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_compressUpdate(ctx[], pointer(buffer), bufsize, pointer(testIn), test_size, C_NULL)
        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_flush(ctx[], pointer(buffer), bufsize, C_NULL)
        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_compressEnd(ctx[], pointer(buffer), bufsize, C_NULL)
    end

    @testset "CompressInvalid" begin
        ctx = Ref{Ptr{CodecLz4.LZ4F_cctx}}(C_NULL)
        CodecLz4.LZ4F_createCompressionContext(ctx, version)

        prefs = Ptr{CodecLz4.LZ4F_preferences_t}(C_NULL)

        bound = CodecLz4.LZ4F_compressBound(test_size, prefs)
        @test bound > 0

        bufsize = bound + CodecLz4.LZ4F_HEADER_SIZE_MAX
        buffer = Vector{UInt8}(undef, ceil(Int, bound / 8))

        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_compressBegin(ctx[], buffer, UInt(2), prefs)
        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_compressUpdate(ctx[], pointer(buffer), bufsize, pointer(testIn), test_size, C_NULL)

        result = CodecLz4.LZ4F_freeCompressionContext(ctx[])
        @test !CodecLz4.LZ4F_isError(result)

        ctx = Ref{Ptr{CodecLz4.LZ4F_cctx}}(C_NULL)
        CodecLz4.LZ4F_createCompressionContext(ctx, version)

        @test_nowarn result = CodecLz4.LZ4F_compressBegin(ctx[], buffer, bufsize, prefs)

        offset = result
        @test_nowarn result = CodecLz4.LZ4F_compressUpdate(ctx[], pointer(buffer) + offset, bufsize - offset, pointer(testIn), test_size, C_NULL)

        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_flush(ctx[], pointer(buffer), UInt(2), C_NULL)
        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_compressEnd(ctx[], pointer(buffer), UInt(2), C_NULL)

        result = CodecLz4.LZ4F_freeCompressionContext(ctx[])
        @test !CodecLz4.LZ4F_isError(result)
    end

    @testset "DecompressUninitialized" begin
        dctx = Ref{Ptr{CodecLz4.LZ4F_dctx}}(C_NULL)
        srcsize = Ref{Csize_t}(test_size)
        dstsize =  Ref{Csize_t}(8*1280)
        decbuffer = Vector{UInt8}(undef, 1280)

        frameinfo = Ref(CodecLz4.LZ4F_frameInfo_t())
        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_getFrameInfo(dctx[], frameinfo, pointer(testIn), srcsize)
        @test_throws CodecLz4.LZ4Exception CodecLz4.LZ4F_decompress(dctx[], decbuffer, dstsize, pointer(testIn), srcsize, C_NULL)
    end

    @testset "Preferences" begin
        ctx = Ref{Ptr{CodecLz4.LZ4F_cctx}}(C_NULL)
        err = CodecLz4.LZ4F_isError(CodecLz4.LZ4F_createCompressionContext(ctx, version))
        @test !err
        opts = Ref(CodecLz4.LZ4F_compressOptions_t(1, (0, 0, 0)))
        prefs = Ref(CodecLz4.LZ4F_preferences_t(CodecLz4.LZ4F_frameInfo_t(), 20, 0, (0, 0, 0, 0)))

        bound = CodecLz4.LZ4F_compressBound(test_size, prefs)
        @test bound > 0

        bufsize = bound + CodecLz4.LZ4F_HEADER_SIZE_MAX
        buffer = Vector{UInt8}(undef, ceil(Int, bound / 8))

        @test_nowarn result = CodecLz4.LZ4F_compressBegin(ctx[], buffer, bufsize, prefs)

        offset = result
        @test_nowarn result = CodecLz4.LZ4F_compressUpdate(ctx[], pointer(buffer) + offset, bufsize - offset, pointer(testIn), test_size, opts)

        offset += result
        @test_nowarn result = CodecLz4.LZ4F_flush(ctx[], pointer(buffer) + offset, bufsize - offset, opts)

        offset += result
        @test_nowarn result = CodecLz4.LZ4F_compressEnd(ctx[], pointer(buffer) + offset, bufsize - offset, opts)
        @test result > 0

        offset += result

        result = CodecLz4.LZ4F_freeCompressionContext(ctx[])
        @test !CodecLz4.LZ4F_isError(result)

        test_decompress(offset, buffer)
    end

end


text:
using PardisoMKL
using Test
using SparseArrays

# Define linear system

mtype = 11

A = sparse([
  0. -2  3 0
  -2  4 -4 1
  -3  5  1 1
  1 -3 0 2 ])

n = A.n

b = Float64[1, 3, 2, 5]

x = zeros(Float64,n)

# Create the pardiso internal handler

pt = new_pardiso_handle()

# pardisoinit!

iparm = new_iparm()

pardisoinit!(pt,mtype,iparm)

# pardiso! (solving the transpose of the system above)

maxfct = 1
mnum = 1
phase = 13
a = A.nzval
ia = Vector{Int32}(A.colptr)
ja = Vector{Int32}(A.rowval)
perm = zeros(Int32,n)
nrhs = 1
msglvl = 0

err = pardiso!(
  pt,maxfct,mnum,mtype,phase,n,a,ia,ja,perm,nrhs,iparm,msglvl,b,x)

tol = 1.0e-13

@test err == 0

@test maximum(abs.(A'*x-b)) < tol

# pardiso! (solving the system above)

iparm[12] = 2

err = pardiso!(
  pt,maxfct,mnum,mtype,phase,n,a,ia,ja,perm,nrhs,iparm,msglvl,b,x)

@test err == 0

@test maximum(abs.(A*x-b)) < tol

@test pardiso_data_type(mtype,iparm) == Float64

# pardiso_getdiag!

iparm[56] = 1

pt = new_pardiso_handle()

err = pardiso!(
  pt,maxfct,mnum,mtype,phase,n,a,ia,ja,perm,nrhs,iparm,msglvl,b,x)

df = zeros(Float64,n)
da = zeros(Float64,n)

err = pardiso_getdiag!(pt,df,da,mnum,mtype,iparm)

@test err == 0

err = pardiso_getdiag!(pt,df,da,mnum)

@test err == 0

# pardiso_64! (solving the transpose of the system above)

pt = new_pardiso_handle()

a = A.nzval
ia = A.colptr
ja = A.rowval
perm = zeros(Int64,n)
iparm = new_iparm_64()

err = pardiso_64!(
  pt,maxfct,mnum,mtype,phase,n,a,ia,ja,perm,nrhs,iparm,msglvl,b,x)

@test err == 0

@test maximum(abs.(A'*x-b)) < tol


text:
using Gtk.ShortNames, Colors

let c = @Canvas(), win = @Window(c, "Canvas1")
    Gtk.draw(c) do widget
        fill!(widget, RGB(1,0,0))
    end
    showall(win)
end

let c = @Canvas(), win = @Window(c, "Canvas2")
    Gtk.draw(c) do widget
        w, h = Int(width(widget)), Int(height(widget))
        randcol = reinterpret(RGB{U8}, rand(0x00:0xff, 3, w, h), (w, h))
        copy!(widget, randcol)
    end
    showall(win)
end


text:
include("header.jl")
include("visual.jl")

opt = ADAMW(Float64(conf["lr"]), (0.9, 0.999), Float64(conf["weight_decay"]));

include("network.jl")

if Bool(conf["is_nn"])
    include("network_nn.jl")
end

include("callback.jl")
include("train.jl")

@suppress_err begin
    include("export.jl")
end


text:
#
# Context awareness.
#

AroundQuery(dom::Domain, refl::Bool, before::Bool, after::Bool) =
    Query(
        AroundSig(refl, before, after, false),
        Input(dom) |> setrelative(),
        Output(dom) |> setoptional(!refl) |> setplural(before || after))

AroundQuery(dom, refl::Bool, before::Bool, after::Bool) =
    AroundQuery(convert(Domain, dom), refl, before, after)

function AroundQuery(dom::Domain, refl::Bool, before::Bool, after::Bool, keys::Vector{Query})
    for key in keys
        @assert fits(dom, domain(input(key)))
        @assert isplain(output(key))
    end
    q = RecordQuery(ItQuery(dom), keys)
    q >>
    Query(
        AroundSig(refl, before, after, true),
        Input(domain(output(q))) |> setrelative(),
        Output(dom) |> setoptional(!refl) |> setplural(before || after))
end

AroundQuery(dom, refl::Bool, before::Bool, after::Bool, keys::Query...) =
    AroundQuery(convert(Domain, dom), refl, before, after, collect(Query, keys))

immutable AroundSig <: AbstractPrimitive
    refl::Bool
    before::Bool
    after::Bool
    haskey::Bool
end

ev(sig::AroundSig, ::Input, oty::Output, iflow::InputFlow) =
    OutputFlow(
        oty,
        if !sig.haskey
            around_impl(sig.refl, sig.before, sig.after, frameoffsets(iflow), values(iflow))
        else
            around_by_impl(sig.refl, sig.before, sig.after, frameoffsets(iflow), values(iflow), fields(domain(iflow)))
        end)

function around_impl(refl::Bool, before::Bool, after::Bool, foffs::AbstractVector{Int}, vals::AbstractVector)
    size = 0
    for k = 1:endof(foffs)-1
        l = foffs[k]
        r = foffs[k+1]
        w = r - l
        if before
            size += w * (w-1) ÷ 2
        end
        if refl
            size += w
        end
        if after
            size += w * (w-1) ÷ 2
        end
    end
    offs = Vector{Int}(length(vals)+1)
    offs[1] = 1
    idxs = Vector{Int}(size)
    n = 1
    for k = 1:endof(foffs)-1
        l = foffs[k]
        r = foffs[k+1]
        for i = l:r-1
            if before
                for j = l:i-1
                    idxs[n] = j
                    n += 1
                end
            end
            if refl
                idxs[n] = i
                n += 1
            end
            if after
                for j = i+1:r-1
                    idxs[n] = j
                    n += 1
                end
            end
            offs[i+1] = n
        end
    end
    return Column{!refl,before||after}(offs, vals[idxs])
end

function around_by_impl(
        refl::Bool, before::Bool, after::Bool, foffs::AbstractVector{Int}, dv::DataVector,
        fs::Vector{Output})
    perm, offs = group_by_keys_impl(foffs, dv, fs)
    rng = Vector{Int}(length(dv))
    pos = Vector{Int}(length(dv))
    size = 0
    for k = 1:endof(offs)-1
        l = offs[k]
        r = offs[k+1]
        for i = l:r-1
            n = perm[i]
            rng[n] = k
            pos[n] = i
        end
        w = r - l
        if before
            size += w * (w-1) ÷ 2
        end
        if refl
            size += w
        end
        if after
            size += w * (w-1) ÷ 2
        end
    end
    idxs = Vector{Int}(size)
    offs′ = Vector{Int}(length(dv)+1)
    offs′[1] = 1
    n = 1
    for k = 1:length(dv)
        r = rng[k]
        p = pos[k]
        l = offs[r]
        r = offs[r+1]
        if before
            for j = l:p-1
                idxs[n] = perm[j]
                n += 1
            end
        end
        if refl
            idxs[n] = perm[p]
            n += 1
        end
        if after
            for j = p+1:r-1
                idxs[n] = perm[j]
                n += 1
            end
        end
        offs′[k+1] = n
    end
    return Column{!refl,before||after}(offs′, values(dv, 1)[idxs])
end

max_stars_repo_name: rbt-lang/rbt-proto

max_issues_repo_name: rbt-lang/rbt-proto

max_forks_repo_name: rbt-lang/rbt-proto


text:
using Darknet
using FileIO
using ImageCore
using Test

datadir = joinpath(dirname(@__DIR__), "data")

@testset "Create Darknet image" begin
    darknet_img = Darknet.array_to_image(rand(Float32,100,10))
    @test typeof(darknet_img) == Darknet.image
    @test darknet_img.h == 100
    @test darknet_img.w == 10
    @test darknet_img.c == 1

    darknet_img = Darknet.array_to_image(rand(Float32,3,100,10))
    @test typeof(darknet_img) == Darknet.image
    @test darknet_img.h == 100
    @test darknet_img.w == 10
    @test darknet_img.c == 3
end

@testset "Read Metadata" begin
    meta = Darknet.get_metadata(joinpath(dirname(@__DIR__),"data","coco.data"))
    @test 80 == meta.classes
end

@testset "Load and run" begin
    weightsfile = "yolov3-tiny.weights"
    cfgfile = "yolov3-tiny.cfg"
    datafile = "coco.data"

    net = Darknet.load_network(joinpath(datadir, cfgfile), joinpath(datadir, weightsfile), 1)
    meta = Darknet.get_metadata(joinpath(datadir, datafile));

    imagefile = joinpath(@__DIR__, "examples", "dog-cycle-car.png")

    img = convert(Array{Float32}, channelview(load(imagefile))) #Read in array via a julia method
    img = img[1:3, :, :] #throw away the alpha channel
    img_d = Darknet.array_to_image(img) #Darknet image type with pointers to source data

    results = Darknet.detect(net, meta, img_d, thresh=0.1, nms=0.3)
    @test length(results) == 7
end


text:
function multiplesimulations(numsim, inputlist... ; rng::AbstractRNG = Random.GLOBAL_RNG)

    multsimlist = [multiplesimulations(numsim, input, rng = rng)
                        for input in inputlist]
    return multsimlist
end

function multiplesimulations(numsim, input; rng::AbstractRNG = Random.GLOBAL_RNG)

    results = MultiSimulation(input, ModuleTracker[])
    for i in 1:numsim
        simdata = run1simulation(input, rng)
        push!(results.output, simdata.output)
    end
    return results
end

function getstats(multsimlist::Array{MultiSimulation{T},1}, args::Symbol...;
    fit_fmin=0.12, fit_fmax=0.24, sampled=true) where T <: SimulationInput
    i = 1
    dflist = DataFrame[]
    for multsim in multsimlist
        μ_predicted = Float64[]
        r2 = Float64[]
        data = sampled ? multsim.sampled : multsim.output
        for simresults in data
            VAF = sampled ? simresults.VAF : simresults.trueVAF
            _, fitcoef, r2val = fitinverse(VAF, fit_fmin, fit_fmax)
            push!(μ_predicted, fitcoef[1])
            push!(r2, r2val)
        end
        μ_error = ((μ_predicted .- multsim.input.μ)
                        ./ multsim.input.μ .* 100)
        df = DataFrame(id = i, μ_predicted = μ_predicted, μ_error = μ_error, r2 = r2)
        addinputparams!(df,multsim,args)
        i += 1
        push!(dflist, df)
    end
    df = reduce(vcat, dflist, cols = :union)
end

function addinputparams!(df,multsim,args)
    for fieldname in args
        if fieldname in fieldnames(InputParameters)
            value = getfield(multsim.input, fieldname)
        elseif fieldname in fieldnames(typeof(multsim.input))
            value = getfield(multsim.input, fieldname)
        end
        addparam!(df,fieldname,value)
    end
end

function addparam!(df,fieldname,value)
    if value isa AbstractArray
        for (i,v) in enumerate(value)
            df[!, Symbol(fieldname, i)] .= v
        end
    else
        df[!, fieldname] .= value
    end
    return df
end


text:
using Test

@testset "Stonks" begin
  include("requests.jl")
  include("apiclients.jl")
  include("parsers_utils.jl")
  include("parsers_yahoo.jl")
  include("parsers_alphavantage.jl")
  include("conversions.jl")
  include("stores_atomic_write.jl")
  include("stores.jl")
  include("stores_updates.jl")
end


text:
const serverCapabilities = ServerCapabilities(
    TextDocumentSyncOptions(true,
    TextDocumentSyncKinds.Full,
    false,
    false,
    SaveOptions(true)),
    CompletionOptions(false, [".", "@", "\"", "^"], missing),
    true,
    SignatureHelpOptions(["(", ","], missing),
    false,
    true,
    false,
    false,
    true,
    false,
    true,
    true,
    missing,
    missing,
    false,
    true,
    false,
    missing,
    true,
    false,
    ExecuteCommandOptions(missing, String["ExplicitPackageVarImport", "ExpandFunction", "AddDefaultConstructor", "ReexportModule", "FixMissingRef"]),
    false,
    true,
    WorkspaceOptions(WorkspaceFoldersOptions(true, true)),
    missing)

hasreadperm(p::String) = (uperm(p) & 0x04) == 0x04

function isjuliabasedir(path)
    try
        fs = readdir(path)
        if "base" in fs && isdir(joinpath(path, "base"))
            return isjuliabasedir(joinpath(path, "base"))
        end
        return all(f -> f in fs, ["coreimg.jl", "coreio.jl", "inference.jl"])
    catch err
        isa(err, Base.IOError) || isa(err, Base.SystemError) || rethrow()
        return false
    end
end

function has_too_many_files(path, N = 5000)
    i = 0

    try
        for (root, dirs, files) in walkdir(path, onerror = x -> x)
            for file in files
                if endswith(file, ".jl")
                    i += 1
                end
                if i > N
                    @info "Your workspace folder has > $N Julia files, server will not try to load them."
                    return true
                end
            end
        end
    catch err
        isa(err, Base.IOError) || isa(err, Base.SystemError) || rethrow()
        return false
    end

    return false
end

function load_rootpath(path)
    try
        return isdir(path) &&
            hasreadperm(path) &&
            path != "" &&
            path != homedir() &&
            !isjuliabasedir(path) &&
            !has_too_many_files(path)
    catch err
        isa(err, Base.IOError) || isa(err, Base.SystemError) || rethrow()
        return false
    end
end

function load_folder(wf::WorkspaceFolder, server)
    path = uri2filepath(wf.uri)
    load_folder(path, server)
end

function load_folder(path::String, server)
    if load_rootpath(path)
        try
            for (root, dirs, files) in walkdir(path, onerror = x -> x)
                for file in files
                    filepath = joinpath(root, file)
                    if isvalidjlfile(filepath)
                        uri = filepath2uri(filepath)
                        if hasdocument(server, URI2(uri))
                            set_is_workspace_file(getdocument(server, URI2(uri)), true)
                            continue
                        else
                            content = try
                                s = read(filepath, String)
                                isvalid(s) || continue
                                s
                            catch err
                                isa(err, Base.IOError) || isa(err, Base.SystemError) || rethrow()
                                continue
                            end
                            doc = Document(uri, content, true, server)
                            setdocument!(server, URI2(uri), doc)
                            parse_all(doc, server)
                        end
                    end
                end
            end
        catch err
            isa(err, Base.IOError) || isa(err, Base.SystemError) || rethrow()
        end
    end
end

function initialize_request(params::InitializeParams, server::LanguageServerInstance, conn)
    # Only look at rootUri and rootPath if the client doesn't support workspaceFolders
    if ismissing(params.capabilities.workspace.workspaceFolders) || params.capabilities.workspace.workspaceFolders == false
        if !(params.rootUri isa Nothing)
            push!(server.workspaceFolders, uri2filepath(params.rootUri))
        elseif !(params.rootPath isa Nothing)
            push!(server.workspaceFolders,  params.rootPath)
        end
    elseif (params.workspaceFolders !== nothing) & (params.workspaceFolders !== missing)
        for wksp in params.workspaceFolders
            if wksp.uri !== nothing
                push!(server.workspaceFolders, uri2filepath(wksp.uri))
            end
        end
    end

    server.clientCapabilities = params.capabilities
    server.clientInfo = params.clientInfo

    if !ismissing(params.capabilities.window) && params.capabilities.window.workDoneProgress
        server.clientcapability_window_workdoneprogress = true
    else
        server.clientcapability_window_workdoneprogress = false
    end

    if !ismissing(params.capabilities.workspace.didChangeConfiguration) &&
        !ismissing(params.capabilities.workspace.didChangeConfiguration.dynamicRegistration) &&
        params.capabilities.workspace.didChangeConfiguration.dynamicRegistration

        server.clientcapability_workspace_didChangeConfiguration = true
    end

    return InitializeResult(serverCapabilities, missing)
end

function initialized_notification(params::InitializedParams, server::LanguageServerInstance, conn)
    server.status = :running

    if server.clientcapability_workspace_didChangeConfiguration
        JSONRPC.send(
            conn,
            client_registerCapability_request_type,
            RegistrationParams([Registration(string(uuid4()), "workspace/didChangeConfiguration", missing)])
        )
    end

    if server.workspaceFolders !== nothing
        for wkspc in server.workspaceFolders
            load_folder(wkspc, server)
        end
    end
    request_julia_config(server, conn)

    if server.number_of_outstanding_symserver_requests > 0
        create_symserver_progress_ui(server)
    end
end

# TODO provide type for params
function shutdown_request(params, server::LanguageServerInstance, conn)
    return nothing
end

# TODO provide type for params
function exit_notification(params, server::LanguageServerInstance, conn)
    server.symbol_server.process isa Base.Process && kill(server.symbol_server.process)
    exit()
end


text:
### A Pluto.jl notebook ###
# v0.12.21

using Markdown
using InteractiveUtils

# ╔═╡ ff8e2e5c-879a-11eb-0d3e-e5049a1cf3b5
begin
  using Pkg; Pkg.activate("MLJ_env", shared=true)
  using PlutoUI

    using CategoricalArrays
    import MLJLinearModels
    import DataFrames
    import CSV
    import DecisionTree
    using MLJ
    import MLJLinearModels
    using Plots
    using Test
end

# ╔═╡ 4f20c7b4-879a-11eb-2209-cf5a5c88437e
md"""
## A workshop introducing the machine learning toolbox MLJ - Part5 - Advanced Model Composition

  - source:[MLJ](https://alan-turing-institute.github.io/MLJ.jl/stable/)
  - based on [Machine Learning in Julia using MLJ, JuliaCon2020](https://github.com/ablaom/MachineLearningInJulia2020)
"""

# ╔═╡ 78069384-879a-11eb-172b-23c56086cf2f
html"""

"""

# ╔═╡ 7940701e-879a-11eb-3753-6d63fa224f3d
md"""
##### TOC
  - [Building a pipeline using the generic composition syntax](#building-pipeline)
  - [A composite model to average two regressor predictors](#composite-model)
  - [Resources or this part](#resources)
"""

# ╔═╡ 79280c84-879a-11eb-27e4-e19e80efbc15
md"""
> **Goals:**
> 1. Learn how to build a prototypes of a composite model, called a *learning network*
> 2. Learn how to use the `@from_network` macro to export a learning network as a new stand-alone model type

While `@pipeline` is great for composing models in an unbranching
sequence, for more complicated model composition you'll want to use
MLJ's generic model composition syntax. There are two main steps:

- **Prototype** the composite model by building a *learning
  network*, which can be tested on some (dummy) data as you build
  it.

- **Export** the learning network as a new stand-alone model type.

Like pipeline models, instances of the exported model type behave
like any other model (and are not bound to any data, until you wrap
them in a machine).
"""

# ╔═╡ 5d151ef4-87a1-11eb-0a27-adf50f61ae71
html"""

  back to TOC

"""

# ╔═╡ 5e014aea-87a1-11eb-2091-833f05e273a5
md"""
#### Building a pipeline using the generic composition syntax

To warm up, we'll do the equivalent of
"""

# ╔═╡ 790e31b2-879a-11eb-0219-8bfd44bbd0a2
pipe = @pipeline Standardizer MLJLinearModels.LogisticClassifier;

# ╔═╡ 78f1e956-879a-11eb-3cb8-cbf7c8c9a2c6
md"""
using the generic syntax.
"""

# ╔═╡ 3038e52e-879b-11eb-155a-3b9240c2e866
md"""
Here's some dummy data we'll be using to test our learning network:
"""

# ╔═╡ 364e5f7a-879b-11eb-2a8b-5b7532fdd2ed
begin
	X₀, y₀ = make_blobs(5, 3)
	with_terminal() do
		pretty(X₀)
	end
end

# ╔═╡ 3ed487b4-879b-11eb-3293-5b2d629353e0
md"""
 **Step 0** - Proceed as if you were combining the models "by hand", using all the data available for training, transforming and prediction:
"""

# ╔═╡ 3ebad7e2-879b-11eb-21f3-f7e3328019d4
begin
	stand₀ = Standardizer();
	linear₀ = MLJLinearModels.LogisticClassifier();
end

# ╔═╡ 3ea2fddc-879b-11eb-0863-479f44d0034b
begin
	mach₀ = machine(stand₀, X₀);
	fit!(mach₀);
	Xstand₀ = transform(mach₀, X₀);
end

# ╔═╡ 3e6e8eaa-879b-11eb-0514-3de8464aec2d
begin
	mach₁ = machine(linear₀, Xstand₀, y₀);
	fit!(mach₁);
	ŷd = predict(mach₁, Xstand₀)
end

# ╔═╡ 1603af44-879c-11eb-034d-43faec4dfa53
md"""
**Step 1** - Edit your code as follows:
  - pre-wrap the data in `Source` nodes
  - delete the `fit!` calls
"""

# ╔═╡ 269c2a48-879c-11eb-062f-c77a84a41250
begin
	X₁ = source(X₀)  # or Xsd = source() if not testing
	y₁ = source(y₀)  # or ysd = source()
end

# ╔═╡ 267f95b8-879c-11eb-3683-4d82a83ce5dd
begin
	stand₁ = Standardizer();
	linear₁ = MLJLinearModels.LogisticClassifier();
end

# ╔═╡ 266347dc-879c-11eb-00d0-579f6bb9d793
begin
	mach₃ = machine(stand₁, X₁);
	Xstand₁ = transform(mach₃, X₁);
end

# ╔═╡ 264c5a72-879c-11eb-3b03-3b935674741d
begin
	mach₄ = machine(linear₁, Xstand₁, y₁);
	ŷ = predict(mach₄, Xstand₁);
end

# ╔═╡ 2615615c-879c-11eb-1ad7-1dca7e88b8b9
md"""
Now `X`, `y`, `Xstand` and `yhat` are *nodes* ("variables" or "dynammic data") instead of data. All training, predicting and transforming is now executed lazily, whenever we `fit!` one of these nodes. $(html"
")
We *call* a node to retrieve the data it represents in the original manual workflow.
"""

# ╔═╡ b9085140-879c-11eb-2bf3-9d7553c6e47d
begin
	fit!(Xstand₁)
	with_terminal() do
		Xstand₁() |> pretty
	end
end

# ╔═╡ b8ea3200-879c-11eb-3109-6564629b90e7
begin
	fit!(ŷ);
	ŷ()
end

# ╔═╡ b8cf3e70-879c-11eb-028e-2fde06bbe2b3
md"""
The node `ŷ` is the "descendant" (in an associated DAG we have defined) of a unique source node:
"""

# ╔═╡ b8b6ee0e-879c-11eb-1855-7d9fd8102ff7
sources(ŷ)

# ╔═╡ b89c6ce6-879c-11eb-0857-3f6d9f8e47d9
md"""
The data at the source node is replaced by `Xnew` to obtain a new prediction when we call `yhat` like this:
"""

# ╔═╡ f08ba086-879c-11eb-12fe-29c5ed220d89
begin
	Xnew, _y = make_blobs(2, 3);
	ŷ(Xnew)
end

# ╔═╡ f0729820-879c-11eb-0791-591900ba2773
md"""
**Step 2** - Export the learning network as a new stand-alone model type

Now, somewhat paradoxically, we can wrap the whole network in a
special machine - called a *learning network machine* - before have
defined the new model type. Indeed doing so is a necessary step in
the export process, for this machine will tell the export macro:

- what kind of model the composite will be (`Deterministic`,
  `Probabilistic` or `Unsupervised`)a

- which source nodes are input nodes and which are for the target

- which nodes correspond to each operation (`predict`, `transform`,
  etc) that we might want to define

"""

# ╔═╡ f05acfbc-879c-11eb-37d5-634f6429e3a5
begin
	surrogate = Probabilistic()     # a model with no fields!
	machₛ = machine(surrogate, X₁, y₁; predict=ŷ);
end

# ╔═╡ f03df408-879c-11eb-3914-b9d1cc896915
md"""
Although we have no real need to use it, this machine behaves like you'd expect it to:
"""

# ╔═╡ f026f51e-879c-11eb-2cbc-7931007bee2c
begin
	Xnew₁, _y₁ = make_blobs(2, 3)
	fit!(machₛ)
	predict(machₛ, Xnew₁)
end

# ╔═╡ a3c94e6e-879d-11eb-36df-6d09e6863d0f
md"""
Now we create a new model type using a Julia `struct` definition appropriately decorated:
"""

# ╔═╡ a3adcef8-879d-11eb-17d7-8303bc978368
@from_network machₛ begin

    mutable struct YaPipe
        standardizer = stand₁
        classifier = linear₁::Probabilistic
    end

end

# ╔═╡ a384ae76-879d-11eb-20df-f55d1cec4ce1
md"""
Instantiating and evaluating on some new data:
"""

# ╔═╡ f00d9e02-879c-11eb-2d82-77383805929e
begin
	pipeᵢ = YaPipe()

	Xᵢ, yᵢ = @load_iris;   # built-in data set

	machᵢ = machine(pipeᵢ, Xᵢ, yᵢ);
	evaluate!(machᵢ, measure=misclassification_rate, operation=predict_mode)
end

# ╔═╡ 8bcd9032-87a1-11eb-2911-8bd0473971bf
html"""

  back to TOC

"""

# ╔═╡ 36d41bbc-879e-11eb-33ec-1d7b44b5abca
md"""
#### A composite model to average two regressor predictors

The following is condensed version of
[this](https://github.com/alan-turing-institute/MLJ.jl/blob/master/binder/MLJ_demo.ipynb)
tutorial. We will define a composite model that:

  - standardizes the input data

  - learns and applies a Box-Cox transformation to the target variable

  - blends the predictions of two supervised learning models - a ridge regressor and a random forest regressor; we'll blend using a simple average (for a more sophisticated stacking example, see [here](https://alan-turing-institute.github.io/DataScienceTutorials.jl/getting-started/stacking/))

  - applies the *inverse* Box-Cox transformation to this blended prediction
"""

# ╔═╡ 36b75e46-879e-11eb-0599-d39bf1a74977
@load RandomForestRegressor pkg=DecisionTree

# ╔═╡ 3599f556-879f-11eb-3f03-0f78dd7b610a
@load RidgeRegressor pkg=MLJLinearModels

# ╔═╡ 79889a18-879f-11eb-0f3a-5d7deb0984a6
@load RandomForestRegressor pkg=DecisionTree

# ╔═╡ 369ce2e4-879e-11eb-26a0-97db50a6048f
md"""
**Input layer**
"""

# ╔═╡ 8fa14d14-879e-11eb-3e7d-4721cce58f3d
begin
	Xₛ = source()
	yₛ = source()
end

# ╔═╡ 8f83c014-879e-11eb-0982-7d918eaeb0e3
md"""
 **First layer and target transformation**
"""

# ╔═╡ a823e9d2-879e-11eb-2cc6-c91241e8ee66
begin
	std_model = Standardizer()
	stand = machine(std_model, Xₛ)
	W = MLJ.transform(stand, Xₛ)
end

# ╔═╡ a80150c0-879e-11eb-022c-a3f5457fd43f
begin
	box_model = UnivariateBoxCoxTransformer()
	box = machine(box_model, yₛ)
	z = MLJ.transform(box, yₛ)
end

# ╔═╡ 8f6ac156-879e-11eb-0ecc-8930e2044e8d
md"""
**Second layer**
"""

# ╔═╡ 8f4ba418-879e-11eb-368b-ddaa316099f0
begin
	ridge_model = MLJLinearModels.RidgeRegressor(lambda=0.1)
	ridge = machine(ridge_model, W, z)
end

# ╔═╡ 3680edc2-879e-11eb-19e6-83bbf95b9a51
begin
	forest_model = MLJDecisionTreeInterface.RandomForestRegressor(n_trees=50)
	forest = machine(forest_model, W, z)
end

# ╔═╡ 36662f9e-879e-11eb-12de-25f09360d3f5
ẑₛ = 0.5 * predict(ridge, W) + 0.5 * predict(forest, W)

# ╔═╡ 9b8a7870-879f-11eb-2671-d9f80b9598d5
md"""
**Output**
"""

# ╔═╡ 9b70e374-879f-11eb-04b5-87a9b8dc7803
ŷₛ = inverse_transform(box, ẑₛ)

# ╔═╡ 9b583018-879f-11eb-181e-91585bf7bdda
md"""
With the learning network defined, we're ready to export:
"""

# ╔═╡ 9b39e4dc-879f-11eb-3ae0-8df5e6a25913
@from_network machine(Deterministic(), Xₛ, yₛ, predict=ŷₛ) begin
    mutable struct CompositeModel
        rgs1 = ridge_model
        rgs2 = forest_model
	end
end

# ╔═╡ 9774caee-87a0-11eb-29d8-f928c2cdc209
md"""
Let's instantiate the new model type and try it out on some data:
"""

# ╔═╡ 97499ccc-87a0-11eb-1465-f1f8a0a2ca75
composite = CompositeModel()

# ╔═╡ 9aedc7fc-879f-11eb-1cc3-af422dcfdc05
begin
	Xₓ, yₓ = @load_boston;

	machₓ = machine(composite, Xₓ, yₓ);
	evaluate!(machₓ,
		resampling=CV(nfolds=6, shuffle=true),
		measures=[rms, mae])
end

# ╔═╡ ed04718c-87a0-11eb-346c-3b459bf27916
html"""

  back to TOC

"""

# ╔═╡ 0a291696-87a1-11eb-1640-f13435ec21ac
md"""
### Resources for Part 5

- From the MLJ manual:
   - [Learning Networks](https://alan-turing-institute.github.io/MLJ.jl/stable/composing_models/#Learning-Networks-1)
- From Data Science Tutorials:
    - [Learning Networks](https://alan-turing-institute.github.io/DataScienceTutorials.jl/getting-started/learning-networks/)
    - [Learning Networks 2](https://alan-turing-institute.github.io/DataScienceTutorials.jl/getting-started/learning-networks-2/)

    - [Stacking](https://alan-turing-institute.github.io/DataScienceTutorials.jl/getting-started/learning-networks-2/)
       an advanced example of model compostion

    - [Finer Control](https://alan-turing-institute.github.io/MLJ.jl/dev/composing_models/#Method-II:-Finer-control-(advanced)-1)
      exporting learning networks without a macro for finer control

"""

# ╔═╡ Cell order:
# ╟─4f20c7b4-879a-11eb-2209-cf5a5c88437e
# ╠═ff8e2e5c-879a-11eb-0d3e-e5049a1cf3b5
# ╟─78069384-879a-11eb-172b-23c56086cf2f
# ╟─7940701e-879a-11eb-3753-6d63fa224f3d
# ╟─79280c84-879a-11eb-27e4-e19e80efbc15
# ╟─5d151ef4-87a1-11eb-0a27-adf50f61ae71
# ╟─5e014aea-87a1-11eb-2091-833f05e273a5
# ╠═790e31b2-879a-11eb-0219-8bfd44bbd0a2
# ╟─78f1e956-879a-11eb-3cb8-cbf7c8c9a2c6
# ╟─3038e52e-879b-11eb-155a-3b9240c2e866
# ╠═364e5f7a-879b-11eb-2a8b-5b7532fdd2ed
# ╟─3ed487b4-879b-11eb-3293-5b2d629353e0
# ╠═3ebad7e2-879b-11eb-21f3-f7e3328019d4
# ╠═3ea2fddc-879b-11eb-0863-479f44d0034b
# ╠═3e6e8eaa-879b-11eb-0514-3de8464aec2d
# ╠═1603af44-879c-11eb-034d-43faec4dfa53
# ╠═269c2a48-879c-11eb-062f-c77a84a41250
# ╠═267f95b8-879c-11eb-3683-4d82a83ce5dd
# ╠═266347dc-879c-11eb-00d0-579f6bb9d793
# ╠═264c5a72-879c-11eb-3b03-3b935674741d
# ╟─2615615c-879c-11eb-1ad7-1dca7e88b8b9
# ╠═b9085140-879c-11eb-2bf3-9d7553c6e47d
# ╠═b8ea3200-879c-11eb-3109-6564629b90e7
# ╟─b8cf3e70-879c-11eb-028e-2fde06bbe2b3
# ╠═b8b6ee0e-879c-11eb-1855-7d9fd8102ff7
# ╟─b89c6ce6-879c-11eb-0857-3f6d9f8e47d9
# ╠═f08ba086-879c-11eb-12fe-29c5ed220d89
# ╟─f0729820-879c-11eb-0791-591900ba2773
# ╠═f05acfbc-879c-11eb-37d5-634f6429e3a5
# ╟─f03df408-879c-11eb-3914-b9d1cc896915
# ╠═f026f51e-879c-11eb-2cbc-7931007bee2c
# ╟─a3c94e6e-879d-11eb-36df-6d09e6863d0f
# ╠═a3adcef8-879d-11eb-17d7-8303bc978368
# ╟─a384ae76-879d-11eb-20df-f55d1cec4ce1
# ╠═f00d9e02-879c-11eb-2d82-77383805929e
# ╟─8bcd9032-87a1-11eb-2911-8bd0473971bf
# ╟─36d41bbc-879e-11eb-33ec-1d7b44b5abca
# ╠═36b75e46-879e-11eb-0599-d39bf1a74977
# ╠═3599f556-879f-11eb-3f03-0f78dd7b610a
# ╠═79889a18-879f-11eb-0f3a-5d7deb0984a6
# ╟─369ce2e4-879e-11eb-26a0-97db50a6048f
# ╠═8fa14d14-879e-11eb-3e7d-4721cce58f3d
# ╟─8f83c014-879e-11eb-0982-7d918eaeb0e3
# ╠═a823e9d2-879e-11eb-2cc6-c91241e8ee66
# ╠═a80150c0-879e-11eb-022c-a3f5457fd43f
# ╟─8f6ac156-879e-11eb-0ecc-8930e2044e8d
# ╠═8f4ba418-879e-11eb-368b-ddaa316099f0
# ╠═3680edc2-879e-11eb-19e6-83bbf95b9a51
# ╠═36662f9e-879e-11eb-12de-25f09360d3f5
# ╟─9b8a7870-879f-11eb-2671-d9f80b9598d5
# ╠═9b70e374-879f-11eb-04b5-87a9b8dc7803
# ╠═9b583018-879f-11eb-181e-91585bf7bdda
# ╠═9b39e4dc-879f-11eb-3ae0-8df5e6a25913
# ╟─9774caee-87a0-11eb-29d8-f928c2cdc209
# ╠═97499ccc-87a0-11eb-1465-f1f8a0a2ca75
# ╠═9aedc7fc-879f-11eb-1cc3-af422dcfdc05
# ╟─ed04718c-87a0-11eb-346c-3b459bf27916
# ╟─0a291696-87a1-11eb-1640-f13435ec21ac

max_stars_repo_name: pascal-p/julia-notebooks

max_issues_repo_name: pascal-p/julia-notebooks

max_forks_repo_name: pascal-p/julia-notebooks


text:
function is_unix(os::Symbol)
    if (os==:Windows) return false;
    elseif (os==:Linux) return true;
    elseif (os==:FreeBSD) return true;
    elseif (os==:Darwin) return true;
    else error("unknown operating system")
    end
end

function _os_test(qm,ex,test)
    @assert qm == :?
    @assert isa(ex,Expr)
    @assert ex.head == :(:)
    @assert length(ex.args) == 2
    if test
        return esc(ex.args[1])
    else
        return esc(ex.args[2])
    end
end
macro windows(qm,ex)
    _os_test(qm, ex, OS_NAME===:Windows)
end
macro unix(qm,ex)
    _os_test(qm, ex, is_unix(OS_NAME))
end
macro osx(qm,ex)
    _os_test(qm, ex, OS_NAME===:Darwin)
end
macro linux(qm,ex)
    _os_test(qm, ex, is_unix(OS_NAME) && OS_NAME!==:Darwin)
end

macro windows_only(ex)
    @windows? esc(ex) : nothing
end
macro unix_only(ex)
    @unix? esc(ex) : nothing
end
macro osx_only(ex)
    @osx? esc(ex) : nothing
end
macro linux_only(ex)
    @linux? esc(ex) : nothing
end

max_stars_repo_name: porterjamesj/julia

max_issues_repo_name: porterjamesj/julia

max_forks_repo_name: porterjamesj/julia


text:
# Defining test functions here as if they are defined where used it is too easy to
# mistakenly create closures over variables that only share names by coincidence.
futestkws(x; err=true) = err ? error("futestkws_err") : x

fbtestkws(x, y; err=true) = err ? error("fbtestkws_err") : x

sinconj(x) = sin(x)

primalapprox(x) = x

f_inferrable(x) = x
f_noninferrable_frule(x) = x
f_noninferrable_rrule(x) = x
f_noninferrable_pullback(x) = x
f_noninferrable_thunk(x, y) = x + y
f_inferrable_pullback_only(x) = x > 0 ? Float64(x) : Float32(x)

function finplace!(x; y=[1])
    y[1] = 2
    x .*= y[1]
    return x
end

struct Foo
    a::Float64
end
(f::Foo)(x) = return f.a + x
Base.length(::Foo) = 1
Base.iterate(f::Foo) = iterate(f.a)
Base.iterate(f::Foo, state) = iterate(f.a, state)

# constructor
function ChainRulesCore.rrule(::Type{Foo}, a)
    foo = Foo(a)
    function Foo_pullback(Δfoo)
        return NoTangent(), unthunk(Δfoo).a
    end
    return foo, Foo_pullback
end
function ChainRulesCore.frule((_, Δa), ::Type{Foo}, a)
    return Foo(a), Foo(Δa)
end

# functor
function ChainRulesCore.rrule(f::Foo, x)
    y = f(x)
    function Foo_pullback(Δy)
        return Tangent{Foo}(;a=Δy), Δy
    end
    return y, Foo_pullback
end
function ChainRulesCore.frule((Δf, Δx), f::Foo, x)
    return f(x), Δf.a + Δx
end

# testing configs
abstract type MySpecialTrait end
struct MySpecialConfig <: RuleConfig{Union{MySpecialTrait}} end

@testset "testers.jl" begin
    @testset "test_scalar" begin
        @testset "Ensure correct rules succeed" begin
            double(x) = 2x
            @scalar_rule(double(x), 2)
            test_scalar(double, 2.1)
        end
        @testset "Ensure incorrect rules caught" begin
            alt_double(x) = 2x
            @scalar_rule(alt_double(x), 3)  # this is wrong, on purpose
            @test fails(() -> test_scalar(alt_double, 2.1))
        end
    end

    @testset "unary: identity(x)" begin
        function ChainRulesCore.frule((_, ẏ), ::typeof(identity), x)
            return x, ẏ
        end
        function ChainRulesCore.rrule(::typeof(identity), x)
            function identity_pullback(ȳ)
                return (NoTangent(), ȳ)
            end
            return x, identity_pullback
        end
        @testset "test_frule" begin
            test_frule(identity, randn())
            test_frule(identity, randn(4))
        end
        @testset "test_rrule" begin
            test_rrule(identity, randn())
            test_rrule(identity, randn(4))
        end
    end

    @testset "Inplace accumulation: identity on Array" begin
        @testset "Correct definitions" begin
            local inplace_used
            function ChainRulesCore.frule((_, ẋ), ::typeof(identity), x::Array)
                ẏ = InplaceableThunk(ȧ -> (inplace_used = true; ȧ .+= ẋ), @thunk(ẋ))
                return identity(x), ẏ
            end
            function ChainRulesCore.rrule(::typeof(identity), x::Array)
                function identity_pullback(ȳ)
                    x̄_ret = InplaceableThunk(@thunk(ȳ)) do ā
                        inplace_used = true
                        ā .+= ȳ
                    end
                    return (NoTangent(), x̄_ret)
                end
                return identity(x), identity_pullback
            end

            inplace_used = false
            test_frule(identity, randn(4))
            @test inplace_used  # make sure we are using, and thus testing the add!

            inplace_used = false
            test_rrule(identity, randn(4))
            @test inplace_used  # make sure we are using, and thus testing the add!
        end

        @testset "Incorrect in-place definitions" begin
            my_identity(value) = value  # we will define bad rules on this
            function ChainRulesCore.frule((_, ẋ), ::typeof(my_identity), x::Array)
                # only the in-place part is incorrect
                ẏ = InplaceableThunk(ȧ -> ȧ .+= 200 .* ẋ, @thunk(ẋ))
                return my_identity(x), ẏ
            end
            function ChainRulesCore.rrule(::typeof(my_identity), x::Array)
                x_dims = size(x)
                function my_identity_pullback(ȳ)
                    # only the in-place part is incorrect
                    x̄_ret = InplaceableThunk(ā -> ā .+= 200 .* ȳ, @thunk(ȳ))
                    return (NoTangent(), x̄_ret)
                end
                return my_identity(x), my_identity_pullback
            end
            @test fails(() -> test_frule(my_identity, randn(4)))
            @test fails(() -> test_rrule(my_identity, randn(4)))
        end
    end

    @testset "check rules are inferrable" begin
        @testset "check inferred" begin
            ChainRulesCore.frule((_, Δx), ::typeof(f_inferrable), x) = (x, Δx)
            function ChainRulesCore.rrule(::typeof(f_inferrable), x)
                f_inferrable_pullback(Δy) = (NoTangent(), Δy)
                return x, f_inferrable_pullback
            end

            test_frule(f_inferrable, 2.0; check_inferred=false)
            test_frule(f_inferrable, 2.0)
            test_rrule(f_inferrable, 2.0; check_inferred=false)
            test_rrule(f_inferrable, 2.0)
            test_scalar(f_inferrable, 2.0)
            test_scalar(f_inferrable, 2.0; check_inferred=false)
        end

        @testset "check not inferred in frule" begin
            function ChainRulesCore.frule((_, Δx), ::typeof(f_noninferrable_frule), x)
                return (x, x > 0 ? Float64(Δx) : Float32(Δx))
            end
            function ChainRulesCore.rrule(::typeof(f_noninferrable_frule), x)
                f_noninferrable_frule_pullback(Δy) = (NoTangent(), Δy)
                return x, f_noninferrable_frule_pullback
            end

            test_frule(f_noninferrable_frule, 2.0; check_inferred=false)
            @test errors(
                () -> test_frule(f_noninferrable_frule, 2.0),
                "does not match inferred return type",
            )

            test_scalar(f_noninferrable_frule, 2.0; check_inferred=false)
            @test errors(
                () -> test_scalar(f_noninferrable_frule, 2.0),
                "does not match inferred return type",
            )
        end

        @testset "check not inferred in rrule" begin
            ChainRulesCore.frule((_, Δx), ::typeof(f_noninferrable_rrule), x) = (x, Δx)
            function ChainRulesCore.rrule(::typeof(f_noninferrable_rrule), x)
                if x > 0
                    f_noninferrable_rrule_pullback(Δy) = (NoTangent(), Δy)
                    return x, f_noninferrable_rrule_pullback
                else
                    return x, _ -> (NoTangent(), Δy) # this is not hit by the used point
                end
            end

            test_rrule(f_noninferrable_rrule, 2.0; check_inferred=false)
            @test errors(
                () -> test_rrule(f_noninferrable_rrule, 2.0),
                "does not match inferred return type",
            )

            test_scalar(f_noninferrable_rrule, 2.0; check_inferred=false)
            @test errors(
                () -> test_scalar(f_noninferrable_rrule, 2.0),
                "does not match inferred return type",
            )
        end

        @testset "check not inferred in pullback" begin
            function ChainRulesCore.rrule(::typeof(f_noninferrable_pullback), x)
                function f_noninferrable_pullback_pullback(Δy)
                    return (NoTangent(), (x > 0 ? Float64 : Float32)(unthunk(Δy)))
                end
                return x, f_noninferrable_pullback_pullback
            end
            test_rrule(f_noninferrable_pullback, 2.0; check_inferred=false)
            @test errors(
                () -> test_rrule(f_noninferrable_pullback, 2.0),
                "does not match inferred return type",
            )
        end

        @testset "check not inferred in thunk" begin
            function ChainRulesCore.rrule(::typeof(f_noninferrable_thunk), x, y)
                function f_noninferrable_thunk_pullback(Δz)
                    ∂x = @thunk(x > 0 ? Float64(unthunk(Δz)) : Float32(unthunk(Δz)))
                    return (NoTangent(), ∂x, Δz)
                end
                return x + y, f_noninferrable_thunk_pullback
            end
            test_rrule(f_noninferrable_thunk, 2.0, 3.0; check_inferred=false)
            @test errors(
                () -> test_rrule(f_noninferrable_thunk, 2.0, 3.0),
                "does not match inferred return type",
            )
        end

        @testset "check non-inferrable primal still passes if pullback inferrable" begin
            function ChainRulesCore.frule((_, Δx), ::typeof(f_inferrable_pullback_only), x)
                T  = x > 0 ? Float64 : Float32
                return T(x), T(Δx)
            end
            function ChainRulesCore.rrule(::typeof(f_inferrable_pullback_only), x)
                function f_inferrable_pullback_only_pullback(Δy)
                    return NoTangent(), oftype(x, unthunk(Δy))
                end
                return x > 0 ? Float64(x) : Float32(x), f_inferrable_pullback_only_pullback
            end
            test_frule(f_inferrable_pullback_only, 2.0; check_inferred=true)
            test_rrule(f_inferrable_pullback_only, 2.0; check_inferred=true)
        end
    end

    @testset "test derivative conjugated in pullback" begin
        ChainRulesCore.frule((_, Δx), ::typeof(sinconj), x) = (sin(x), cos(x) * Δx)

        # define rrule using ChainRulesCore's v0.9.0 convention, conjugating the derivative
        # in the rrule
        function ChainRulesCore.rrule(::typeof(sinconj), x)
            sinconj_pullback(ΔΩ) = (NoTangent(), conj(cos(x)) * ΔΩ)
            return sin(x), sinconj_pullback
        end

        test_rrule(sinconj, randn(ComplexF64))
        test_scalar(sinconj, randn(ComplexF64))
    end

    @testset "binary: fst(x, y)" begin
        fst(x, y) = x
        ChainRulesCore.frule((_, dx, dy), ::typeof(fst), x, y) = (x, dx)
        function ChainRulesCore.rrule(::typeof(fst), x, y)
            function fst_pullback(Δx)
                return (NoTangent(), Δx, ZeroTangent())
            end
            return x, fst_pullback
        end
        @testset "test_frule" begin
            test_frule(fst, 2.0, 3.0)
            test_frule(fst, randn(4), randn(4))
        end
        @testset "test_rrule" begin
            test_rrule(fst, 2.0, 4.0)
            test_rrule(fst, randn(4), randn(4))
        end
    end

    @testset "single input, multiple output" begin
        simo(x) = (x, 2x)
        function ChainRulesCore.rrule(simo, x)
            simo_pullback((a, b)) = (NoTangent(), a .+ 2 .* b)
            return simo(x), simo_pullback
        end
        function ChainRulesCore.frule((_, ẋ), simo, x)
            y = simo(x)
            return y, Tangent{typeof(y)}(ẋ, 2ẋ)
        end

        @testset "test_frule" begin
            test_frule(simo, 1.5)  # on scalar
            test_frule(simo, randn(4))  # on array
        end
        @testset "test_rrule" begin
            test_rrule(simo, 3.0)  # on scalar
            test_rrule(simo, randn(4))  # on array
        end
    end

    @testset "tuple input: first" begin
        ChainRulesCore.frule((_, dx), ::typeof(first), xs::Tuple) = (first(xs), first(dx))
        function ChainRulesCore.rrule(::typeof(first), x::Tuple)
            function first_pullback(Δx)
                return (NoTangent(), Tangent{typeof(x)}(Δx, falses(length(x) - 1)...))
            end
            return first(x), first_pullback
        end

        @testset "test_frule" begin
            test_frule(first, (2.0, 3.0))
            test_frule(first, Tuple(randn(4)))
        end
        @testset "test_rrule" begin
            test_rrule(first, (2.0, 3.0); check_inferred=false)
            test_rrule(first, Tuple(randn(4)); check_inferred=false)
        end
    end

    @testset "tuple output (backing type of Tangent =/= natural differential)" begin
        tuple_out(x) = return (x, 1.0) # i.e. (x, 1.0) and not (x, x)
        function ChainRulesCore.frule((_, dx), ::typeof(tuple_out), x)
            Ω = tuple_out(x)
            ∂Ω = Tangent{typeof(Ω)}(dx, ZeroTangent())
            return Ω, ∂Ω
        end
        test_frule(tuple_out, 2.0)
    end

    @testset "ignoring arguments" begin
        fsymtest(x, s::Symbol) = x
        ChainRulesCore.frule((_, Δx, _), ::typeof(fsymtest), x, s) = (x, Δx)
        function ChainRulesCore.rrule(::typeof(fsymtest), x, s)
            function fsymtest_pullback(Δx)
                return NoTangent(), Δx, NoTangent()
            end
            return x, fsymtest_pullback
        end

        @testset "test_frule" begin
            test_frule(fsymtest, 2.5, :x)
        end

        @testset "test_rrule" begin
            test_rrule(fsymtest, 2.5, :x)
        end
    end

    @testset "unary with kwargs: futestkws(x; err)" begin
        function ChainRulesCore.frule((_, ẋ), ::typeof(futestkws), x; err=true)
            return futestkws(x; err=err), ẋ
        end
        function ChainRulesCore.rrule(::typeof(futestkws), x; err=true)
            function futestkws_pullback(Δx)
                return (NoTangent(), Δx)
            end
            return futestkws(x; err=err), futestkws_pullback
        end

        # we defined these functions at top of file to throw errors unless we pass `err=false`
        @test_throws ErrorException futestkws(randn())
        @test errors(() -> test_scalar(futestkws, randn()), "futestkws_err")
        @test_throws ErrorException frule((nothing, randn()), futestkws, randn())
        @test_throws ErrorException rrule(futestkws, randn())

        @test_throws ErrorException futestkws(randn(4))
        @test_throws ErrorException frule((nothing, randn(4)), futestkws, randn(4))
        @test_throws ErrorException rrule(futestkws, randn(4))

        @testset "scalar_test" begin
            test_scalar(futestkws, 2.5; fkwargs=(; err=false))
        end
        @testset "test_frule" begin
            test_frule(futestkws, 2.5; fkwargs=(; err=false))
            test_frule(futestkws, 2.5; fkwargs=(; err=false))
        end
        @testset "test_rrule" begin
            test_rrule(futestkws, 2.5; fkwargs=(; err=false))
            test_rrule(futestkws, 2.5; fkwargs=(; err=false))
        end
    end

    @testset "binary with kwargs: fbtestkws(x, y; err)" begin
        function ChainRulesCore.frule((_, ẋ, _), ::typeof(fbtestkws), x, y; err=true)
            return fbtestkws(x, y; err=err), ẋ
        end
        function ChainRulesCore.rrule(::typeof(fbtestkws), x, y; err=true)
            function fbtestkws_pullback(Δx)
                return (NoTangent(), Δx, ZeroTangent())
            end
            return fbtestkws(x, y; err=err), fbtestkws_pullback
        end

        # we defined these functions at top of file to throw errors unless we pass `err=false`
        @test_throws ErrorException fbtestkws(randn(), randn())
        @test_throws ErrorException frule(
            (nothing, randn(), nothing), fbtestkws, randn(), randn()
        )
        @test_throws ErrorException rrule(fbtestkws, randn(), randn())

        @test_throws ErrorException fbtestkws(randn(4), randn(4))
        @test_throws ErrorException frule(
            (nothing, randn(4), nothing), fbtestkws, randn(4), randn(4)
        )
        @test_throws ErrorException rrule(fbtestkws, randn(4), randn(4))

        @testset "test_frule" begin
            test_frule(fbtestkws, 2.5, 3.0; fkwargs=(; err=false))
            test_frule(fbtestkws, randn(4), randn(4); fkwargs=(; err=false))
        end
        @testset "test_rrule" begin
            test_rrule(fbtestkws, 2.5, 3.0; fkwargs=(; err=false))
            test_rrule(fbtestkws, randn(4), randn(4); fkwargs=(; err=false))
        end
    end

    @testset "primal can be only approximately equal" begin
        ChainRulesCore.frule((_, Δx), ::typeof(primalapprox), x) = (x + sqrt(eps(x)), Δx)

        function ChainRulesCore.rrule(::typeof(primalapprox), x)
            function primalapprox_pullback(Δx)
                return (NoTangent(), Δx)
            end
            return x + sqrt(eps(x)), primalapprox_pullback
        end

        test_frule(primalapprox, 2.5; atol=1e-6)
        test_rrule(primalapprox, 2.5; atol=1e-6)
    end

    @testset "frule with mutation" begin
        function ChainRulesCore.frule((_, ẋ), ::typeof(finplace!), x; y=[1])
            y[1] *= 2
            x .*= y[1]
            ẋ .*= 2 # hardcoded to match y defined below
            return x, ẋ
        end

        # these pass in tangents explictly so that we can check them after
        x = randn(3)
        ẋ = [4.0, 5.0, 6.0]
        xcopy, ẋcopy = copy(x), copy(ẋ)
        y = [1, 2]
        test_frule(finplace!, x ⊢ ẋ; fkwargs=(y=y,))
        @test x == xcopy
        @test ẋ == ẋcopy
        @test y == [1, 2]
    end

    @testset "TestIterator input" begin
        function iterfun(iter)
            state = iterate(iter)
            state === nothing && error()
            (x, i) = state
            s = x^2
            while true
                state = iterate(iter, i)
                state === nothing && break
                (x, i) = state
                s += x^2
            end
            return s
        end

        function ChainRulesCore.frule((_, Δiter), ::typeof(iterfun), iter)
            iter_Δiter = zip(iter, Δiter)
            state = iterate(iter_Δiter)
            state === nothing && error()
            # for some reason the following line errors if the frule is defined within a testset
            ((x, Δx), i) = state
            return iterfun(iter), sum(2 .* iter.data .* Δiter.data)
            s = x^2
            ∂s = 2 * x * Δx
            while true
                state = iterate(iter_Δiter, i)
                state === nothing && break
                ((x, Δx), i) = state
                s += x^2
                ∂s += 2 * x * Δx
            end
            return s, ∂s
        end

        function ChainRulesCore.rrule(::typeof(iterfun), iter::TestIterator)
            function iterfun_pullback(Δs)
                data = iter.data
                ∂data = (2 * Δs) .* conj.(data)
                ∂iter = TestIterator(
                    ∂data, Base.IteratorSize(iter), Base.IteratorEltype(iter)
                )
                return (NoTangent(), ∂iter)
            end
            return iterfun(iter), iterfun_pullback
        end

        @testset "Testing iterator function" begin
            # This needs to be in a separate testset to stop the `x` being shared with `iterfun`
            x = TestIterator(randn(2, 3), Base.SizeUnknown(), Base.EltypeUnknown())
            test_frule(iterfun, x)
            test_rrule(iterfun, x)
        end
    end

    @testset "unhappy path" begin
        @testset "primal wrong" begin
            my_identity1(x) = x
            function ChainRulesCore.frule((_, ẏ), ::typeof(my_identity1), x)
                return 2.5 * x, ẏ
            end
            function ChainRulesCore.rrule(::typeof(my_identity1), x)
                function identity_pullback(ȳ)
                    return (NoTangent(), ȳ)
                end
                return 2.5 * x, identity_pullback
            end

            @test fails(() -> test_frule(my_identity1, 2.2))
            @test fails(() -> test_rrule(my_identity1, 2.2))
        end

        @testset "derivative wrong" begin
            my_identity2(x) = x
            function ChainRulesCore.frule((_, ẏ), ::typeof(my_identity2), x)
                return x, 2.7 * ẏ
            end
            function ChainRulesCore.rrule(::typeof(my_identity2), x)
                function identity_pullback(ȳ)
                    return (NoTangent(), 31.8 * ȳ)
                end
                return x, identity_pullback
            end

            @test fails(() -> test_frule(my_identity2, 2.2))
            @test fails(() -> test_rrule(my_identity2, 2.2))
        end

        @testset "wrong number of outputs #167" begin
            foo(x, y) = x + 2y

            function ChainRulesCore.frule((_, ẋ, ẏ), ::typeof(foo), x, y)
                return foo(x, y), ẋ + 2ẏ, NoTangent() # extra derivative
            end

            function ChainRulesCore.rrule(::typeof(foo), x, y)
                foo_pullback(dz) = NoTangent(), dz # missing derivative
                return foo(x,y), foo_pullback
            end

            @test fails(() -> test_frule(foo, 2.1, 2.1))
            @test fails(() -> test_rrule(foo, 21.0, 32.0))
        end

        @testset "rrule not returning a tuple" begin
            bar(x, y) = x + 3y
            function ChainRulesCore.rrule(::typeof(bar), x, y)
                bar_pullback(dy) = dy
                return bar(x,y), bar_pullback
            end
            @test fails(() -> test_rrule(bar, 21.0, 32.0))
        end
    end

    @testset "structs" begin
        @testset "constructor" begin
            test_frule(Foo, rand())
            test_rrule(Foo, rand())
        end

        foo = Foo(rand())
        tfoo = Tangent{Foo}(;a=rand())
        @testset "functor" begin
            test_frule(foo, rand())
            test_rrule(foo, rand())
            test_scalar(foo, rand())

            test_frule(foo ⊢ Foo(rand()), rand())
            test_frule(foo ⊢ tfoo, rand())
            test_rrule(foo ⊢ Foo(rand()), rand())
            test_rrule(foo ⊢ tfoo, rand())
        end
    end

    @testset "Tuple primal that is not equal to differential backing" begin
        # https://github.com/JuliaMath/SpecialFunctions.jl/issues/288
        forwards_trouble(x) = (1, 2.0 * x)
        @scalar_rule(forwards_trouble(v), ZeroTangent(), 2.0)
        test_frule(forwards_trouble, 2.5)

        rev_trouble((x, y)) = y
        function ChainRulesCore.rrule(::typeof(rev_trouble), (x, y)::P) where {P}
            rev_trouble_pullback(ȳ) = (NoTangent(), Tangent{P}(ZeroTangent(), ȳ))
            return y, rev_trouble_pullback
        end
        test_rrule(rev_trouble, (3, 3.0) ⊢ Tangent{Tuple{Int,Float64}}(ZeroTangent(), 1.0))
    end

    @testset "check_thunked_output_tangent" begin
        @testset "no method for thunk" begin
            does_not_accept_thunk_id(x) = x
            function ChainRulesCore.rrule(::typeof(does_not_accept_thunk_id), x)
                does_not_accept_thunk_id_pullback(ȳ::AbstractArray) = (NoTangent() ,ȳ)
                return does_not_accept_thunk_id(x), does_not_accept_thunk_id_pullback
            end

            test_rrule(
                does_not_accept_thunk_id, [1.0, 2.0]; check_thunked_output_tangent=false
            )
            @test errors(r"MethodError.*Thunk") do
                test_rrule(does_not_accept_thunk_id, [1.0, 2.0])
            end
        end

        @testset "Thunk wrong" begin
            bad_thunk_id(x) = x
            function ChainRulesCore.rrule(::typeof(bad_thunk_id), x)
                bad_thunk_id_pullback(ȳ::AbstractArray) = (NoTangent(), ȳ)
                bad_thunk_id_pullback(ȳ::AbstractThunk) = (NoTangent(), 2 * ȳ)
                return bad_thunk_id(x), bad_thunk_id_pullback
            end

            test_rrule(bad_thunk_id, [1.0, 2.0]; check_thunked_output_tangent=false)
            @test fails(()->test_rrule(bad_thunk_id, [1.0, 2.0]))
        end
    end

    @testset "error message about incorrectly using ZeroTangent()" begin
        foo(a, i) = a[i]
        function ChainRulesCore.rrule(::typeof(foo), a, i)
            function foo_pullback(Δy)
                da = zeros(size(a))
                da[i] = Δy
                return NoTangent(), da, ZeroTangent()
            end
            return foo(a, i), foo_pullback
        end
        @test errors(() -> test_rrule(foo, [1.0, 2.0, 3.0], 2), "should use NoTangent()")
    end

    @testset "NotImplemented" begin
        f_notimplemented(x, y) = (x + y, x - y)
        @scalar_rule f_notimplemented(x, y) (@not_implemented(""), 1) (1, -1)
        test_frule(f_notimplemented, randn(), randn())
        test_rrule(f_notimplemented, randn(), randn())
    end

    @testset "custom rrule_f" begin
        only2x(x, y) = 2x
        custom(::RuleConfig, ::typeof(only2x), x, y) = only2x(x, y), Δ -> (NoTangent(), 2Δ, ZeroTangent())
        wrong1(::RuleConfig, ::typeof(only2x), x, y) = only2x(x, y), Δ -> (ZeroTangent(), 2Δ, ZeroTangent())
        wrong2(::RuleConfig, ::typeof(only2x), x, y) = only2x(x, y), Δ -> (NoTangent(), 2.1Δ, ZeroTangent())
        wrong3(::RuleConfig, ::typeof(only2x), x, y) = only2x(x, y), Δ -> (NoTangent(), 2Δ)

        test_rrule(only2x, 2.0, 3.0; rrule_f=custom, check_inferred=false)
        @test errors(() -> test_rrule(only2x, 2.0, 3.0; rrule_f=wrong1, check_inferred=false))
        @test fails(() -> test_rrule(only2x, 2.0, 3.0; rrule_f=wrong2, check_inferred=false))
        @test fails(() -> test_rrule(only2x, 2.0, 3.0; rrule_f=wrong3, check_inferred=false))
    end

    @testset "custom frule_f" begin
        mytuple(x, y) = return 2x, 1.0
        T = Tuple{Float64, Float64}
        custom(::RuleConfig, (Δf, Δx, Δy), ::typeof(mytuple), x, y) = mytuple(x, y), Tangent{T}(2Δx, ZeroTangent())
        wrong1(::RuleConfig, (Δf, Δx, Δy), ::typeof(mytuple), x, y) = mytuple(x, y), Tangent{T}(2.1Δx, ZeroTangent())
        wrong2(::RuleConfig, (Δf, Δx, Δy), ::typeof(mytuple), x, y) = mytuple(x, y), Tangent{T}(2Δx, 1.0)

        test_frule(mytuple, 2.0, 3.0; frule_f=custom, check_inferred=false)
        @test fails(() -> test_frule(mytuple, 2.0, 3.0; frule_f=wrong1, check_inferred=false))
        @test fails(() -> test_frule(mytuple, 2.0, 3.0; frule_f=wrong2, check_inferred=false))
    end

    @testset "custom_config" begin
        has_config(x) = 2x
        function ChainRulesCore.rrule(::MySpecialConfig, ::typeof(has_config), x)
            has_config_pullback(ȳ) = return (NoTangent(), 2ȳ)
            return has_config(x), has_config_pullback
        end

        has_trait(x) = 2x
        function ChainRulesCore.rrule(::RuleConfig{<:MySpecialTrait}, ::typeof(has_trait), x)
            has_trait_pullback(ȳ) = return (NoTangent(), 2ȳ)
            return has_trait(x), has_trait_pullback
        end

        # it works if the special config is provided
        test_rrule(MySpecialConfig(), has_config, rand())
        test_rrule(MySpecialConfig(), has_trait, rand())

        # but it doesn't work for the default config
        errors(() -> test_rrule(has_config, rand()), "no method matching rrule")
        errors(() -> test_rrule(has_trait, rand()), "no method matching rrule")
    end

    @testset "@maybe_inferred" begin
        f_noninferrable(x) = Ref{Real}(x)[]

        @test @maybe_inferred(identity(1)) == 1
        @test errors("return type $Int does not match inferred return type Real") do
            @maybe_inferred f_noninferrable(1)
        end
        @test @maybe_inferred(Real, f_noninferrable(1)) == 1

        ChainRulesTestUtils.TEST_INFERRED[] = false

        @test @maybe_inferred(identity(1)) == 1
        @test @maybe_inferred(f_noninferrable(1)) == 1
        @test @maybe_inferred(Real, f_noninferrable(1)) == 1

        ChainRulesTestUtils.TEST_INFERRED[] = true
    end

    @testset "inference of thunked cotangents" begin
        my_id(x) = x
        function ChainRulesCore.rrule(::typeof(my_id), x)
            my_id_pb(ȳ) = (NoTangent(), ȳ)
            function my_id_pb(ȳ::AbstractThunk)
                precision = rand() > 0.5 ? Float64 : Float32
                return (NoTangent(), precision(unthunk(ȳ)))
            end
            return x, my_id_pb
        end

        @test errors(() -> test_rrule(my_id, 2.0))
        test_rrule(my_id, 2.0; check_inferred=false)
        test_rrule(my_id, 2.0; check_thunked_output_tangent=false)
    end
end


text:
type mVertex
    index::Int64
    value::Any    # will be of type string
end

type mEdge
    index::Int64
    v1::mVertex
    v2::mVertex
    value::Any
end

""" define a simple graph type with vertices and edges that have unique indices
# as well as attributes.
# Also define an out_edges adjacency list
# and an in_edges adjacency list """
type Multigraph
    vertices::Dict{Int64,mVertex}
    edges::Dict{Int64,mEdge}
    out_edges::Dict{Int64,Array{Int64}}
    in_edges::Dict{Int64,Array{Int64}}
end

Multigraph() = Multigraph(Dict(),Dict(),Dict(),Dict())

function add_vertex!(g::Multigraphs.Multigraph, next_ind::Int, value)
    v = mVertex(next_ind, value)
    g.vertices[next_ind] = v # add to vertex list
    g.out_edges[next_ind] = [] # push empty adjacency lists
    g.in_edges[next_ind] = []
    return v
end

function add_vertex!(g::Multigraphs.Multigraph, value)
    next_ind = length(g.vertices) + 1
    return add_vertex!(g, next_ind, value)
end

function add_edge!(g::Multigraphs.Multigraph, v1_ind::Int, v2_ind::Int, value)
    edge_ind = length(g.edges) + 1
    v1 = g.vertices[v1_ind]
    v2 = g.vertices[v2_ind]
    e = mEdge(edge_ind, v1, v2, value)
    g.edges[edge_ind] = e
    push!(g.out_edges[v1_ind], edge_ind)
    push!(g.in_edges[v2_ind], edge_ind)
    return e
end

function remove_edge!(g::Multigraphs.Multigraph, edge::Multigraphs.mEdge)
    delete!(g.edges, edge.index)
    filter!(x->x!=edge.index, g.out_edges[edge.v1.index])
    filter!(x->x!=edge.index, g.in_edges[edge.v2.index])
    return g
end

function add_edge!(g::Multigraphs.Multigraph, v1::mVertex, v2::mVertex, value)
    return add_edge!(g, v1.index, v2.index, value)
end

in_degree(g::Multigraphs.Multigraph, v::Multigraphs.mVertex) = length(g.in_edges[v.index])

out_degree(g::Multigraphs.Multigraph,v::Multigraphs.mVertex) = length(g.out_edges[v.index])


text:
using Revise
using SnpArrays
using DelimitedFiles
using Random
using MendelIHT
using Knockoffs
using LinearAlgebra
using GLMNet
using Distributions
using DataFrames
using CSV
using Printf
BLAS.set_num_threads(1)

# predict with estimated β̂ (R2 = 1 - RSS/TSS)
function R2(X::AbstractMatrix, y::AbstractVector, β̂::AbstractVector)
    μ = y - X * β̂
    tss = y .- mean(y)
    return 1 - dot(μ, μ) / dot(tss, tss)
end

# predict with a low dimensional fit
function R2(Xtrain::AbstractMatrix, Xtest::AbstractMatrix,
    ytrain::AbstractVector, ytest::AbstractVector, β̂::AbstractVector)
    # fit low dimensional model on original data
    idx = findall(!iszero, β̂)
    β_new = zeros(length(β̂))
    β_new[idx] .= Xtrain[:, idx] \ ytrain
    # predict with low diemensional model on new data
    μ = ytest - Xtest * β_new
    t = ytest .- mean(ytest)
    return 1 - dot(μ, μ) / dot(t, t)
end

function TP(correct_groups, signif_groups)
    return length(signif_groups ∩ correct_groups) / length(correct_groups)
end

function FDR(correct_groups, signif_groups)
    FP = length(signif_groups) - length(signif_groups ∩ correct_groups) # number of false positives
    # FPR = FP / (FP + TN) # https://en.wikipedia.org/wiki/False_positive_rate#Definition
    FDR = FP / length(signif_groups)
    return length(signif_groups) == 0 ? 0 : FDR
end

function tune_k(y::AbstractVector, xko_la::AbstractMatrix, covar::AbstractVecOrMat,
    original::Vector{Int}, knockoff::Vector{Int}, fdr::Float64, best_k::Int
    )
    # do a grid search for best sparsity level
    best_result = nothing
    best_err = Inf
    for cur_k in best_k:5:round(Int, 2best_k)
        result = fit_iht(y, xko_la, covar, k=cur_k, init_beta=true, max_iter=500)
        W = coefficient_diff(result.beta, original, knockoff)
        τ = threshold(W, fdr, :knockoff)
        detected = count(x -> x ≥ τ, W)
        if abs(detected - best_k) < best_err
            best_result = deepcopy(result)
            best_err = abs(detected - best_k)
        end
        println("wrapped CV says best_k = $best_k; using k = $cur_k detected $detected")
        GC.gc()
    end
    return best_result
end
# function tune_k(y::AbstractVector, xko_la::AbstractMatrix, original::Vector{Int},
#     knockoff::Vector{Int}, fdr::Float64, best_k::Int
#     )
#     cur_k = best_k
#     detected = 0
#     for iter in 1:5
#         result = fit_iht(y, xko_la, k=cur_k, init_beta=true, max_iter=500)
#         W = coefficient_diff(result.beta, original, knockoff)
#         τ = threshold(W, fdr, :knockoff)
#         detected = count(x -> x ≥ τ, W)
#         if detected < best_k
#             cur_k += best_k - detected
#         else
#             break
#         end
#         GC.gc()
#     end
#     @show result
#     println("wrapped CV says best_k = $best_k; after tuning detected = $detected")
#     return result.beta
# end

# todo: bolt
function run_sims(seed::Int;
    k = 10, # number of causal SNPs
    use_PCA = false,
    combine_beta=false,
    extra_k = 0,
    confounders = 0, # number of confounders
    causal_snp_upper_r2 = 1, # causal SNPs and their knockoffs must have correlation less than causal_snp_upper_r2
    causal_snp_lower_r2 = 0  # causal SNPs and their knockoffs must have correlation larger than causal_snp_lower_r2
    )
    #
    # import fastphase knockoffs
    #
    # chr = 10
    # plinkname = "/scratch/users/bbchu/ukb/subset/ukb.10k.chr$chr"
    # original = vec(readdlm("/scratch/users/bbchu/ukb/subset/ukb.chr$chr.original.snp.index", Int))
    # knockoff = vec(readdlm("/scratch/users/bbchu/ukb/subset/ukb.chr$chr.knockoff.snp.index", Int))
    # x = SnpArray(plinkname)
    # xko = SnpArray(knockoffname)
    # xla = convert(Matrix{Float64}, x, center=true, scale=true, impute=true)
    # xko_la = convert(Matrix{Float64}, xko, center=true, scale=true, impute=true)
    # cur_dir = pwd()

    #
    # import shapeit knockoffs
    #
    chr = 10
    plinkname = "/scratch/users/bbchu/ukb_SHAPEIT/knockoffs/ukb_gen_chr$(chr)_ibd1_res0" #shapeit knockoffs
    # plinkname = "/scratch/users/bbchu/ukb_SHAPEIT/decorrelated_knockoffs/ukb_gen_chr$(chr)_ibd1_res0_decorrelated" #shapeit knockoffs (decorrelated)
    xdata = SnpData(plinkname)
    isknockoff = endswith.(xdata.snp_info[!, :snpid], ".k")
    original, knockoff = Int[], Int[]
    for i in 1:size(xdata)[2]
        isknockoff[i] ? push!(knockoff, i) : push!(original, i)
    end
    x = xdata.snparray
    xla = convert(Matrix{Float64}, @view(x[1:10000, original]), center=true, scale=true, impute=true)
    xko_la = convert(Matrix{Float64}, @view(x[1:10000, :]), center=true, scale=true, impute=true)
    cur_dir = pwd()

    #
    # Import PCs
    # if use_PCA, make augmented design matrix for lasso and covarirate matrix for IHT
    #
    z = readdlm("/scratch/users/bbchu/ukb_SHAPEIT/subset_pca/ukb.10k.chr$chr.projections.txt")
    standardize!(z)
    xla_full = use_PCA ? [xla z] : xla
    xko_la_full = use_PCA ? [xko_la z] : xko_la
    covar = use_PCA ? [ones(size(xla, 1)) z] : ones(size(xla, 1))

    #
    # simulate phenotypes using UKB chr10 subset
    #
    n, p = size(xla)
    # simulate β
    Random.seed!(seed)
    h2 = 0.5 # heritability
    d = Normal(0, sqrt(h2 / (2k))) # from paper: Efficient Implementation of Penalized Regression for Genetic Risk Prediction
    β = zeros(p)
    # β[1:k] .= rand(d, k)
    # shuffle!(β)
    # causal SNPs can only be those that aren't very correlated with their knockoffs
    possible_causal_snp_idx = Int[]
    for snp in 1:size(xla, 2)
        r2 = cor(@view(xko_la[:, 2snp]), @view(xko_la[:, 2snp - 1]))
        causal_snp_lower_r2 ≤ r2 ≤ causal_snp_upper_r2 && push!(possible_causal_snp_idx, snp)
    end
    shuffle!(possible_causal_snp_idx)
    β[possible_causal_snp_idx[1:k]] .= rand(d, k)
    # simulate y
    ϵ = Normal(0, 1 - h2)
    y = xla * β + rand(ϵ, n)
    #
    # confounders are PCs and have effect size ±0.2
    #
    if confounders > 0
        PCs = z[:, 1:confounders]
        γ = [rand(-1:2:1) * 0.2 for i in 1:confounders]
        y += PCs * γ
    end

    #
    # Run standard IHT
    #
    Random.seed!(seed)
    path = 10:10:200
    mses = cv_iht(y, xla, covar, path=path, init_beta=true)
    GC.gc()
    Random.seed!(seed)
    k_rough_guess = path[argmin(mses)]
    dense_path = (k_rough_guess - 9):(k_rough_guess + 9)
    mses_new = cv_iht(y, xla, covar, path=dense_path, init_beta=true)
    GC.gc()
    Random.seed!(seed)
    iht_result = fit_iht(y, xla, covar, k=dense_path[argmin(mses_new)], init_beta=true, max_iter=500)
    @show iht_result
    GC.gc()

    #
    # Run standard lasso
    #
    Random.seed!(seed)
    @time lasso_cv = glmnetcv(xla_full, y, nfolds=5, parallel=true)

    #
    # run knockoff IHT
    #
    Random.seed!(seed)
    path = 10:10:200
    mses = cv_iht(y, xko_la, covar, path=path, init_beta=true)
    GC.gc()
    Random.seed!(seed)
    k_rough_guess = path[argmin(mses)]
    dense_path = (k_rough_guess - 9):(k_rough_guess + 9)
    mses_new = cv_iht(y, xko_la, covar, path=dense_path, init_beta=true)
    GC.gc()
    Random.seed!(seed)
    iht_ko_result = fit_iht(y, xko_la, covar, k=dense_path[argmin(mses_new)]+extra_k,
        init_beta=true, max_iter=500)
    @show iht_ko_result

    #
    # Run knockoff lasso
    #
    Random.seed!(seed)
    @time lasso_ko_cv = glmnetcv(xko_la_full, y, nfolds=5, parallel=true)

    #
    # Run bolt lmm (need to make covariate and phenotype file first)
    #
    # bedfile = data_dir * "train"
    # covfile = "cov_bolt.txt"
    # phefile = "y_bolt.txt"
    # outfile = "bolt_output.txt"
    # open(phefile, "w") do io
    #     println(io, "FID IID trait1")
    #     for i in 1:size(xla, 1)
    #         println(io, "$i 1 $(ytrain[i])")
    #     end
    # end
    # if use_PCA
    #     open(covfile, "w") do io
    #         println(io, "FID IID PC1 PC2 PC3 PC4 PC5")
    #         for i in 1:size(xla, 1)
    #             println(io, "$i 1 ", z[i, 1], ' ', z[i, 2], ' ',  z[i, 3], ' ', z[i, 4], ' ', z[i, 5])
    #         end
    #     end
    #     run(`$bolt_exe --bfile=$bedfile --covarFile=$covfile
    #         --phenoFile=$phefile --phenoCol=trait1
    #         --qCovarCol=PC\{1:5\} --lmmInfOnly --numLeaveOutChunks=2 --statsFile $outfile`)
    # else
    #     run(`$bolt_exe --bfile=$bedfile --phenoFile=$phefile --phenoCol=trait1
    #         --lmmInfOnly --numLeaveOutChunks=2 --statsFile $outfile`)
    # end

    for fdr in [0.05, 0.1, 0.25, 0.5]
        top_dir = cur_dir * "/fdr$fdr/"
        new_dir = cur_dir * "/fdr$fdr/sim$seed"
        isdir(top_dir) || mkdir(top_dir)
        isdir(new_dir) || mkdir(new_dir)
        cd(new_dir)
        writedlm("y_true.txt", y)
        writedlm("beta_true.txt", β)

        # save IHT, lasso, iht+knockoff, lasso+knockoff results
        writedlm("iht.beta", iht_result.beta)
        writedlm("iht.covariates", iht_result.c)
        writedlm("lasso.beta", coef(lasso_cv))
        writedlm("iht.knockoff.beta", iht_ko_result.beta)
        writedlm("iht.knockoff.covariates", iht_ko_result.c)
        writedlm("lasso.knockoff.beta", coef(lasso_ko_cv))

        #
        # run knockoff IHT with wrapped cross validation
        #
        Random.seed!(seed)
        chr = 10
        path = 10:10:200
        mses = cv_iht_knockoff(y, xko_la, covar, original, knockoff, fdr, path=path,
            init_beta=true, combine_beta = combine_beta)
        Random.seed!(seed)
        GC.gc()
        k_rough_guess = path[argmin(mses)]
        dense_path = (k_rough_guess - 9):(k_rough_guess + 9)
        mses_new = cv_iht_knockoff(y, xko_la, covar, original, knockoff, fdr,
            path=dense_path, init_beta=true, combine_beta = combine_beta)
        GC.gc()
        # adjust sparsity level so it best matches sparsity chosen by ko filter
        Random.seed!(seed)
        best_k = dense_path[argmin(mses_new)]
        best_result = tune_k(y, xko_la, covar, original, knockoff, fdr, best_k)
        writedlm("iht.knockoff.cv.beta", best_result.beta)
        writedlm("iht.knockoff.cv.covariates", best_result.c)

        #
        # compare R2 across populations, save result in a dataframe
        #
        # combine_beta = false
        p = length(β)
        # β_iht = β_iht_knockoff = β_iht_knockoff_cv = zeros(p)
        β_iht = iht_result.beta
        β_lasso = coef(lasso_cv)[1:p]
        β_iht_knockoff = extract_beta(iht_ko_result.beta, fdr,
            original, knockoff, :knockoff, combine_beta)
        β_iht_knockoff_cv = extract_beta(best_result.beta,
            fdr, original, knockoff, :knockoff, combine_beta)
        β_lasso_knockoff = extract_beta(coef(lasso_ko_cv), fdr,
            original, knockoff, :knockoff, combine_beta)[1:p]

        writedlm("iht.knockoff.beta.postfilter", β_iht_knockoff)
        writedlm("iht.knockoff.cv.beta.postfilter", β_iht_knockoff_cv)
        writedlm("lasso.knockoff.beta.postfilter", β_lasso_knockoff)

        populations = ["african", "asian", "bangladeshi", "british", "caribbean", "chinese",
            "indian", "irish", "pakistani", "white_asian", "white_black", "white"]

        df = DataFrame(pop = String[], IHT_R2 = Float64[], IHT_ko_R2 = Float64[],
            IHT_ko_cv_R2 = Float64[], LASSO_R2 = Float64[], LASSO_ko_R2 = Float64[])

        for pop in populations
            xtest = SnpArray("/scratch/users/bbchu/ukb_SHAPEIT/populations/chr10/ukb.chr$chr.$pop.bed") # 10k samples with all snps
            Xtest = SnpLinAlg{Float64}(xtest, center=true, scale=true, impute=true)
            # simulate "true" phenotypes for these populations
            Random.seed!(seed)
            ytest = Xtest * β + rand(ϵ, size(Xtest, 1))
            # IHT
            iht_r2 = R2(Xtest, ytest, β_iht)
            # IHT knockoff (low dimensional fit)
            iht_ko_r2 = R2(xla, Xtest, y, ytest, β_iht_knockoff)
            # knockoff IHT cv (low dimensional)
            iht_ko_cv_r2 = R2(xla, Xtest, y, ytest, β_iht_knockoff_cv)
            # lasso β
            lasso_r2 = R2(Xtest, ytest, β_lasso)
            # knockoff lasso β (low dimensional)
            lasso_ko_r2 = R2(xla, Xtest, y, ytest, β_lasso_knockoff)
            # save to dataframe
            push!(df, hcat(pop, iht_r2, iht_ko_r2, iht_ko_cv_r2,
                lasso_r2, lasso_ko_r2))
            # push!(df, hcat(pop, 0, 0, 0, lasso_r2, lasso_ko_r2))
            GC.gc()
        end

        # count non-zero entries of β returned from cross validation
        push!(df, hcat("beta_non_zero_count", count(!iszero, β_iht),
            count(!iszero, iht_ko_result.beta),
            count(!iszero, best_result.beta),
            # 0, 0,
            count(!iszero, β_lasso),
            count(!iszero, coef(lasso_ko_cv))
            ))
        # count non-zero entries after knockoff filter
        push!(df, hcat("beta_selected", count(!iszero, β_iht),
            count(!iszero, β_iht_knockoff),
            count(!iszero, β_iht_knockoff_cv), count(!iszero, β_lasso),
            count(!iszero, β_lasso_knockoff)
            ))
        # count TP proportion
        correct_snps = findall(!iszero, β)
        push!(df, hcat("TPP", TP(correct_snps, findall(!iszero, β_iht)),
            TP(correct_snps, findall(!iszero, β_iht_knockoff)),
            TP(correct_snps, findall(!iszero, β_iht_knockoff_cv)),
            TP(correct_snps, findall(!iszero, β_lasso)),
            TP(correct_snps, findall(!iszero, β_lasso_knockoff))
            ))
        # count FDR
        push!(df, hcat("FDR", FDR(correct_snps, findall(!iszero, β_iht)),
            FDR(correct_snps, findall(!iszero, β_iht_knockoff)),
            FDR(correct_snps, findall(!iszero, β_iht_knockoff_cv)),
            FDR(correct_snps, findall(!iszero, β_lasso)),
            FDR(correct_snps, findall(!iszero, β_lasso_knockoff))
            ))

        @show df
        CSV.write("summary.txt", df)
    end
end

#
# Run simulation (via `julia prs.jl n`)
# where n is a seed
#
seed = parse(Int, ARGS[1])
# causal_snp_upper_r2 = parse(Float64, ARGS[2]) # 0.1, 0.2, ..., 1.0
# causal_snp_lower_r2 = causal_snp_upper_r2 - 0.1 # 0, 0.1, ..., 0.9
k = 100
confounders = 0 # 1 PC
run_sims(seed, k=k, use_PCA=false, confounders=confounders)