text:
using NFFT, DataFrames, LinearAlgebra, LaTeXStrings, DelimitedFiles
using BenchmarkTools
using Plots, StatsPlots, CategoricalArrays
pgfplotsx()
#gr()

include("../../Wrappers/NFFT3.jl")
include("../../Wrappers/FINUFFT.jl")
mkpath("./img/")
mkpath("./data/")

const threads = [1,2,4,8]
const precomp = [NFFT.POLYNOMIAL, NFFT.TENSOR, NFFT.POLYNOMIAL, NFFT.TENSOR]
const packagesCtor = [NFFTPlan, NFFTPlan, FINUFFTPlan, NFFT3Plan]
const packagesStr = ["NFFT.jl/POLY", "NFFT.jl/TENSOR", "FINUFFT", "NFFT3"]
const benchmarkTime = [1, 20, 20]
#const NBase = [4*4096, 256, 32]
const NBase = [512*512, 512, 64]

NFFT.FFTW.set_num_threads(Threads.nthreads())
ccall(("omp_set_num_threads",NFFT3.lib_path_nfft),Nothing,(Int64,),convert(Int64,Threads.nthreads()))
@info ccall(("nfft_get_num_threads",NFFT3.lib_path_nfft),Int64,())
NFFT._use_threads[] = (Threads.nthreads() > 1)

function nfft_performance_comparison(m = 4, σ = 2.0)
  println("\n\n ##### nfft_performance ##### \n\n")

  df = DataFrame(Package=String[], Threads=Int[], D=Int[], M=Int[], N=Int[],
                   Undersampled=Bool[], m = Int[], σ=Float64[],
                   TimePre=Float64[], TimeTrafo=Float64[], TimeAdjoint=Float64[] )

  fftflags = NFFT.FFTW.MEASURE

  for D = 2:2
      NN = ntuple(d->NBase[D], D)
      M = prod(NN) #÷ 8

        @info D, NN, M

        T = Float64

        x = T.(rand(T,D,M) .- 0.5)

        fHat = randn(Complex{T}, M)
        f = randn(Complex{T}, NN)

        for pl = 1:length(packagesStr)

          planner = packagesCtor[pl]
          BenchmarkTools.DEFAULT_PARAMETERS.seconds = benchmarkTime[1]
          b = @benchmark $planner($x, $NN; m=$m, σ=$σ, window=:kaiser_bessel,
                                 precompute=$(precomp[pl]), sortNodes=false, fftflags=$fftflags)
          tpre = minimum(b).time / 1e9
          p = planner(x, NN; m=m, σ=σ, window=:kaiser_bessel,
                      precompute=(precomp[pl]), sortNodes=false, fftflags=fftflags)
          BenchmarkTools.DEFAULT_PARAMETERS.seconds = benchmarkTime[2]
          b = @benchmark mul!($f, $(adjoint(p)), $fHat)
          tadjoint = minimum(b).time / 1e9
          BenchmarkTools.DEFAULT_PARAMETERS.seconds = benchmarkTime[3]
          b = @benchmark mul!($fHat, $p, $f)
          ttrafo = minimum(b).time / 1e9

          push!(df, (packagesStr[pl], Threads.nthreads(), D, M, NBase[D], false, m, σ,
                   tpre, ttrafo, tadjoint))

    end
  end
  return df
end

function plot_performance(df; D=2, N=1024, M=N*N)

  Plots.scalefontsizes()
  Plots.scalefontsizes(1.5)

  labelsA = packagesStr

  tpre = zeros(length(threads),length(labelsA))
  ttrafo = zeros(length(threads), length(labelsA))
  tadjoint = zeros(length(threads), length(labelsA))
  for (i,p) in enumerate(packagesStr)
    for (j,th) in enumerate(threads) #.&& df.Pre.==precomp[i]
      df_ = df[df.Threads .== th .&& df.Package.==p .&& df.D .== D .&& df.M .== M .&& df.N .==N ,:]
      tpre[j,i] = df_[1,:TimePre]
      ttrafo[j,i] = df_[1,:TimeTrafo]
      tadjoint[j,i] = df_[1,:TimeAdjoint]
    end
  end

  labelsB = "t = " .* string.(threads)

  ctg = CategoricalArray(repeat(labelsA, inner = length(threads)))
  levels!(ctg, labelsA)
  name = CategoricalArray(repeat(labelsB, outer = length(labelsA)))
  levels!(name, labelsB)

  maxtime = max(maximum(tpre), maximum(ttrafo), maximum(tadjoint))

  p1 = groupedbar(name, ttrafo, ylabel = "time / s",  group = ctg,
          bar_width = 0.67, legend = :none, ylims=(0,maxtime),
          lw = 0, framestyle = :box, size=(800,600), title = L"\textrm{NFFT}")

  p2 = groupedbar(name, tadjoint, ylabel = "time / s",  group = ctg,
          bar_width = 0.67, legend = :none, ylims=(0,maxtime),
          lw = 0, framestyle = :box, size=(800,600), title = L"\textrm{NFFT}^H")

  p3 = groupedbar(name, tpre, ylabel = "time / s",  group = ctg,
          bar_width = 0.67, ylims=(0,maxtime),
          lw = 0, framestyle = :box, size=(800,600), title = L"\textrm{Precompute}")

  p = plot(p1, p2, p3, layout=(3,1), size=(800,600), dpi=200)

  savefig(p, "./img/performance_mt.pdf")
  return p
end

function plot_performance_speedup(df; D=2, N=1024, M=N*N)

  Plots.scalefontsizes()
  Plots.scalefontsizes(1.5)

  labelsA = packagesStr

  tpre = zeros(length(threads),length(labelsA))
  ttrafo = zeros(length(threads), length(labelsA))
  tadjoint = zeros(length(threads), length(labelsA))
  for (i,p) in enumerate(packagesStr)
    for (j,th) in enumerate(threads) #.&& df.Pre.==precomp[i]
      df_ = df[df.Threads .== th .&& df.Package.==p .&& df.D .== D .&& df.M .== M .&& df.N .==N ,:]
      tpre[j,i] = df_[1,:TimePre]
      ttrafo[j,i] = df_[1,:TimeTrafo]
      tadjoint[j,i] = df_[1,:TimeAdjoint]
    end
  end

  efftrafo = 1 ./ ttrafo .* ttrafo[1:1,:] ./ threads
  effadjoint = 1 ./ tadjoint .* tadjoint[1:1,:] ./ threads

  ttrafo = 1 ./ ttrafo * ttrafo[1,3]
  tadjoint = 1 ./ tadjoint * tadjoint[1,3]

  Plots.scalefontsizes()
  Plots.scalefontsizes(1.5)

  colors = [RGB(0.0,0.29,0.57), RGB(0.3,0.5,0.7), RGB(0.94,0.53,0.12), RGB(0.99,0.75,0.05)]
  ls = [:solid, :solid, :solid, :solid]
  shape = [:xcross, :circle, :xcross, :cross]

    titleTrafo = L"\textrm{NFFT}"
    titleAdjoint = L"\textrm{NFFT}^H"

    p1 = plot(threads,
              ttrafo[:,1], ylims=(0.0,maximum(threads)),
              label=packagesStr[1], lw=2, ylabel="Speedup", #xlabel = "# threads",
              legend = :topleft, title=titleTrafo,
              shape=shape[1], ls=ls[1],
              c=colors[1], msc=colors[1], mc=colors[1], ms=4, msw=2)

    for p=2:length(packagesStr)
      plot!(p1, threads,
             ttrafo[:,p],
              label=packagesStr[p], lw=2, shape=shape[p], ls=ls[p],
              c=colors[p], msc=colors[p], mc=colors[p], ms=4, msw=2)
    end

    p2 = plot(threads, tadjoint[:,1],  ylims=(0.0,maximum(threads)),
              lw=2, #xlabel = "# threads",
              label=packagesStr[1],
              #legend = i==2 ? :topright : nothing,
              legend = nothing,
              title=titleAdjoint,
              shape=shape[1], ls=ls[1],
              c=colors[1], msc=colors[1], mc=colors[1], ms=4, msw=2)

    for p=2:length(packagesStr)
      plot!(p2, threads, tadjoint[:,p],
              label=packagesStr[p], lw=2, shape=shape[p], ls=ls[p],
              c=colors[p], msc=colors[p], mc=colors[p], ms=4, msw=2)
    end

   p = plot(p1, p2, layout=(1,2), size=(800,300), dpi=200)

  p3 = plot(threads,
              efftrafo[:,1], ylims=(0.0,1.2),
              label=packagesStr[1], lw=2, ylabel="Efficiency", xlabel = "# threads",
              legend = nothing, title=titleTrafo,
              shape=shape[1], ls=ls[1],
              c=colors[1], msc=colors[1], mc=colors[1], ms=4, msw=2)

    for p=2:length(packagesStr)
      plot!(p3, threads,
             efftrafo[:,p],
              label=packagesStr[p], lw=2, shape=shape[p], ls=ls[p],
              c=colors[p], msc=colors[p], mc=colors[p], ms=4, msw=2)
    end

    p4 = plot(threads, effadjoint[:,1],  ylims=(0.0,1.2),
              lw=2, xlabel = "# threads", label=packagesStr[1],
              #legend = i==2 ? :topright : nothing,
              legend = nothing,
              title=titleAdjoint,
              shape=shape[1], ls=ls[1],
              c=colors[1], msc=colors[1], mc=colors[1], ms=4, msw=2)

    for p=2:length(packagesStr)
      plot!(p4, threads, effadjoint[:,p],
              label=packagesStr[p], lw=2, shape=shape[p], ls=ls[p],
              c=colors[p], msc=colors[p], mc=colors[p], ms=4, msw=2)
    end

   p = plot(p1, p2, p3, p4, layout=(2,2), size=(800,500), dpi=200)

  savefig(p, "./img/performance_mt_speedup.pdf")

  return
end

### run the code ###

if haskey(ENV, "NFFT_PERF")
  df = nfft_performance_comparison(4, 2.0)

  if isfile("./data/performance_mt.csv")
    data, header = readdlm("./data/performance_mt.csv", ',', header=true);
    df_ = DataFrame(data, vec(header))
    append!(df, df_)
  end

  writedlm("./data/performance_mt.csv", Iterators.flatten(([names(df)], eachrow(df))), ',')

else
  if false
    rm("./data/performance_mt.csv", force=true)
    ENV["NFFT_PERF"] = 1
    for t in threads
      cmd = `julia -t $t performanceThreading.jl`
      @info cmd
      run(cmd)

    end
    delete!(ENV, "NFFT_PERF")
  end
  data, header = readdlm("./data/performance_mt.csv", ',', header=true);
  df = DataFrame(data, vec(header))

  plot_performance(df, N=NBase[2], M=NBase[2]*NBase[2])
  plot_performance_speedup(df, N=NBase[2], M=NBase[2]*NBase[2])
end


text:
using Base.Iterators: partition
using Flux
using Flux.Optimise: update!
using Flux.Losses: mse, mae
using Images
using Statistics
using Parameters: @with_kw
using Random
using Printf
using Zygote
using CycleGAN: Discriminator, Generator

size_img = 28

@with_kw struct HyperParams
    batch_size::Int = 5
    epochs::Int = 100
    verbose_freq::Int = 800
    size_dataset::Int = 1000
    lr_dscr_A::Float64 = 0.00005
    lr_gen_A::Float64 = 0.00005
    lr_dscr_B::Float64 = 0.00005
    lr_gen_B::Float64 = 0.00005
end

function convertI2Float(img)
    img_resize = imresize(float.(img), (size_img,size_img))
    if length(size(img_resize)) == 2
        img_resize = RGB.(img_resize)
    end
	return permutedims(channelview(img_resize), (3,2,1))
end

function load_images(path::String, size::Int)
	images= zeros(Float32,size_img,size_img,3,size)
	for (index, img) in enumerate(readdir(path, join = true))
		images[:,:,:,index] = convertI2Float(load(img))
        if index == size
            break
        end
	end
	return images
end

function load_data(hparams)
    # Load folder dataset
    images_A = load_images("horse2zebra/trainA/", hparams.size_dataset)
    images_B = load_images("horse2zebra/trainB/", hparams.size_dataset)
    data = [ (images_A[:,:,:, r], images_B[:,:,:, r]) |> gpu for r in partition(1:hparams.size_dataset, hparams.batch_size)]
    return data
end

function create_output_image(gen, image)
    @eval Flux.istraining() = false
    fake_image = cpu(gen(image))
    @eval Flux.istraining() = true
    image_array = permutedims(dropdims(fake_image; dims=4), (3,2,1))
    image_array = colorview(RGB, image_array)
    return clamp01nan.(image_array)
end

function train()
    hparams = HyperParams()

    data = load_data(hparams)

    #test images
    test_images_A=zeros(Float32,size_img,size_img,3,1)
    test_images_B=zeros(Float32,size_img,size_img,3,1)
    test_images_A[:,:,:,1] = convertI2Float(load("horse2zebra/testA/n02381460_1000.jpg"))
    test_images_B[:,:,:,1] = convertI2Float(load("horse2zebra/testB/n02391049_100.jpg"))

    # Discriminator
    dscr_A = Discriminator()
    dscr_B = Discriminator()

    # Generator
    gen_A =  Generator(3, 64) |> gpu
    gen_B =  Generator(3, 64) |> gpu

    # Optimizers
    opt_dscr = ADAM(hparams.lr_dscr_A, (0.5,0.99))
    opt_gen = ADAM(hparams.lr_gen_A, (0.5,0.99))

    cyclegGAN = Cyclegan(gen_A, gen_B, dscr_A, dscr_B)

    isdir("output")||mkdir("output")

    # Training
    train_steps = 0
    for ep in 1:hparams.epochs
        @info "Epoch $ep"
        for (x,y) in data
                # Update discriminator and generator
            loss = train_gan(cyclegGAN, x, y, opt_gen, opt_dscr)
            if train_steps % hparams.verbose_freq == 0
                @info("Train step $(ep), Discriminator loss = $(loss["D_loss"]), Generator loss = $(loss["G_loss"])")
                # Save generated fake image
                output_image_A = create_output_image(cyclegGAN.gen_AB, test_images_B |> gpu)
                output_image_B = create_output_image(cyclegGAN.gen_BA, test_images_A |> gpu)
                save(@sprintf("output/cgan_A_steps_%06d.png", train_steps), output_image_A)
                save(@sprintf("output/cgan_B_steps_%06d.png", train_steps), output_image_B)
            end
        train_steps += 1
        end
    end
    @info("Finish  training")
    output_image_A = create_output_image(cyclegGAN.gen_AB, test_images_B |> gpu)
    output_image_B = create_output_image(cyclegGAN.gen_BA, test_images_A |> gpu)
    save("output/cgan_A_steps_final.png", output_image_A)
    save("output/cgan_B_steps_final.png", output_image_B)
end

cd(@__DIR__)
train()

ext: jl


text:
using Documenter
using AngularMomentumAlgebra
using AtomicLevels
using HalfIntegers
using LinearAlgebra

DocMeta.setdocmeta!(AngularMomentumAlgebra, :DocTestSetup,
                    :(using AngularMomentumAlgebra, AtomicLevels, HalfIntegers, LinearAlgebra); recursive=true)

makedocs(
    modules = [AngularMomentumAlgebra],
    format = Documenter.HTML(assets = ["assets/latex.js"],
                             mathengine = Documenter.MathJax()),
    sitename = "AngularMomentumAlgebra",
    authors="Stefanos Carlström <stefanos.carlstrom@gmail.com>",
    pages = [
        "Home" => "index.md",
        "Definitions" => "definitions.md",
        "Common routines" => "common.md",
        "Orbitals" => "orbitals.md",
        "Tensors" => [
            "General tensors" => "tensors.md",
            "Tensor DSL" => "tensor_dsl.md",
            "Various tensors" => [
                "Angular momenta" => "angular_momenta.md",
                "Spherical tensors" => "spherical_tensors.md",
                "Gradients" => "gradients.md",
                "Multipole expansions" => "multipole_expansions.md",
                "Coulomb interaction" => "coulomb.md",
            ],
            "Tensor matrix elements" => "tensor_matrix_elements.md",
        ],
        "Energy expressions" => "energy_expressions.md",
    ],
    repo="https://github.com/JuliaAtoms/AngularMomentumalgebra.jl/blob/{commit}{path}#L{line}",
)

deploydocs(repo = "github.com/JuliaAtoms/AngularMomentumAlgebra.jl.git")


text:
import Base: *, sort
using Random, Statistics, DataFrames, CSV
using BioSequences,RCall,ArgParse
import BioSequences: iscompatible

################################################
#OligoRL Toolset
#################################################

# Overload the concatenation operator to combine a sequence and
# a single base, i.e. dna"AGCGTGC" * DNA_T
function *(seq::LongDNASeq, base::DNA)
    seq * LongDNASeq([base])
end

# Overload sort to order DNA codes by degeneracy:
#   A,G,C,T,M,R,W,S,Y,K,V,H,D,B,N
function sort(seq::LongSequence{DNAAlphabet{4}}; rev::Bool=false)
    sort(convert(Vector{DNA}, seq), by=degeneracy, rev=rev)
end

# Determine if a given sequence is palindromic. Works both for unambiguous and ambigous bases. Returns TRUE if the sequence is a palindrome.

"""
    is_palindrome(seq)

Check if seq is palindromic.
"""
function is_palindrome(seq)
    return occursin(seq,reverse_complement(seq))
end

"""
    degeneracy(base::DNA[, uselog2=true])

Calculate the degeneracy of a single base. The degeneracy
is the number of bases in the code, so
    degeneracy(DNA_N) == 4
and
    degeneracy(DNA_A) == 1

If `uselog2`, the degeneracy is log2 transformed. By convention,
the degneracy of `DNA_Gap` is zero in both the log2-transformed
and the untransformed case.
"""
function degeneracy(base::DNA; uselog2=true)
    deg = 0.0
    if base == DNA_N
        deg = 4.0
    elseif base == DNA_A || base == DNA_C || base == DNA_G || base == DNA_T
        deg = 1.0
    elseif base == DNA_B || base == DNA_D || base == DNA_H || base == DNA_V
        deg = 3.0
    elseif (base == DNA_K || base == DNA_Y || base == DNA_S || base == DNA_W
            || base == DNA_R || base == DNA_M)
        deg = 2.0
    else # DNA_Gap
        deg = 0.0
    end

    if uselog2
        if base == DNA_Gap
            deg = 2^-100  #log2(0) is negative infinity, so we use something close to zero
        end

        return log2(deg)

    else
        return deg
    end
end

"""
    degeneracy(seq::LongSequence{DNAAlphabet{4}}; uselog2=true)

Calculate the degeneracy of a DNA sequence.
"""
function degeneracy(seq::LongSequence{DNAAlphabet{4}}; uselog2=true)
    if uselog2
        return sum(degeneracy.(seq, uselog2=true))
    else
        return prod(degeneracy.(seq, uselog2=false))
    end
end

#Calculate the degeneracy of an entire oligo pool

"""
    pooled_degeneracy(pool)

Calcualte the sum degeneracy of pool.
"""
function pooled_degeneracy(pool)
    deg=0;
    for randomer in pool
        deg +=degeneracy(randomer;uselog2=false);
    end
    return deg
end

#Determine if two oligos contain at least one common sequence. Used for assessing ambigous oligos.
#Extends the iscompatible function in the BioSequences package from bases to sequences.
#The two sequences must be the same length, otherwise the function returns false

"""
    iscompatible(seq_A::LongSequence{DNAAlphabet{4}},seq_B::LongSequence{DNAAlphabet{4}})

Check if seq_A and seq_B contain at least one common sequence. Seq_A and Seq_B must be the same length.
"""
function iscompatible(seq_A::LongSequence{DNAAlphabet{4}},seq_B::LongSequence{DNAAlphabet{4}})
    if length(seq_A) == length(seq_B)
        return prod(iscompatible.(seq_A,seq_B))
    else # Sequences are different lenghts
        return false
    end
end

#Decompress an ambiguous oligo into all of its possible sequences. Input is a single oligo.
#
"""
    decompress_oligo(comp_oligo)

Expand a degnerate oligo into the non-degnerate oligos it defines.
"""

function decompress_oligo(compressed_oligo::LongSequence{DNAAlphabet{4}})
    basedict=Dict(DNA_A => dna"A", DNA_T => dna"T", DNA_C => dna"C", DNA_G => dna"G",DNA_R =>dna"AG",DNA_Y=>dna"CT",DNA_S=>dna"GC",
    DNA_W=>dna"AT",DNA_K=>dna"GT",DNA_M=>dna"AC",DNA_B=>dna"CGT",DNA_D=>dna"AGT",DNA_H=>dna"ACT",DNA_V=>dna"ACG",DNA_N=>dna"ACGT",DNA_Gap=>dna"-")
    nseqs=trunc(Int128,degeneracy(compressed_oligo;uselog2=false))
    randomerlen=length(compressed_oligo)
    blockarray=zeros(Int128,randomerlen+1)
    blockarray[1]=nseqs
    degsarray=zeros(Int128,randomerlen)
    repsarray=zeros(Int128,randomerlen+1)
    repsarray[1]=1
    uncompressed_oligos=[dna"-"^randomerlen for _ in 1:nseqs]
    for i = 1:randomerlen
        nucleotides=basedict[compressed_oligo[i]]
        degsarray[i]=length(nucleotides)
        blockarray[i+1]=trunc(Int128,blockarray[i]/degsarray[i]);
        bases=dna"-"^blockarray[i];
        repsarray[i+1]=trunc(Int128,prod(degsarray[1:i]));
            for  k in 1:degsarray[i]
                for l in 1:blockarray[i+1]
                    bases[l+blockarray[i+1]*(k-1)]=nucleotides[k]
                end
            end

        bases=bases^repsarray[i];
        for m in 1:nseqs
            uncompressed_oligos[m][i]=bases[m];
        end
    end
    return uncompressed_oligos
end

"""
    decompress_pool(compressed_pool)

Expand a pool of degenerate oligos into a single pool of non-degnerate oligos.
"""
function decompress_pool(compressed_pool)
    decompressed_pool=decompress_oligo(compressed_pool[1])
    for i = 2:length(compressed_pool)
        subset=decompress_oligo(compressed_pool[i])
        for i in subset
            push!(decompressed_pool,i)
        end
    end
    return decompressed_pool
end

#Generate a random pool without replacement. Specify the randomer length, the size of the pool, and the available nucleotides.
#It can create ambiguous sequences, but it is possible that unique sequences will appear multiple times.

"""
    generate_random_pool(;randomerlen=6,randpoolsize=100,nucleotides=dna"ACGT")

Generate a random pool of oligos with length=randomerlen and pool size=randpoolsize using nucleotides
"""
function generate_random_pool(;randomerlen=6,randpoolsize=100,nucleotides=dna"ACGT")
    randpool=[]
    randomer=dna"-"^randomerlen
    for i = 1:randomerlen
        bases=shuffle(nucleotides)
        randomer[i]=bases[1];
    end
    push!(randpool,randomer)
    while length(randpool)<randpoolsize
        randomer=dna"-"^randomerlen
        for i = 1:randomerlen
            bases=shuffle(nucleotides)
            randomer[i]=bases[1];
        end
        check=true
        for j = 1:length(randpool)
            if randomer==randpool[j]
                check=false
                break
            end
        end
        if check
            push!(randpool,randomer)
        end
        if length(randpool)==length(nucleotides)^randomerlen
            break
        end
    end

    return randpool
end

# Given a pool of sequences all of length n, find all of the unambiguous sequences of length n that don't appear in that pool.

"""
    valid_candidates(sites)

Find all non-degenerate oligos that don't appear in sites. All oligos in sites must be of equal length.
"""
function valid_candidates(sites)
    lengths=length.(sites)
    if all(lengths[1].==lengths)
        n=length(sites[1])
        nsites=length(sites)
        candidates=generate_random_pool(;randomerlen=n,randpoolsize=4^n);
        for i=1:nsites
            candidates=candidates[findall(j -> !iscompatible(j,sites[i]),candidates)]
        end
        return candidates
    else
        error("provide an array of oligos of the same length")
    end
end

# Given a list of valid candidates, determine the action space of those candidates. For each base position, return a list of compatible nucleotides.
# If for example, the first postion does not contain an A or a T in any of the valid candidates, then the action space for that base position is dna"GCS".
#returns bases, an array of dna sequences containing all of the valid actions for the agent to take at each position.

"""
    define_action_space(candidates)

Find the available base codes for an oligo, such that only codes that appear in candidates at each postion are included.
"""
function define_action_space(candidates)
    if length(candidates)==0
        bases=[dna"" for i in 1:6]
        return bases
    else
        randomerlen=length(candidates[1]);
        ncands=length(candidates)
        bases=[];
        for i=1:randomerlen
            states=dna"-"^ncands;
            nucleotides=dna"----" # A list of nucleotides not present at a given position in the oligos, called "states" here
            #println("TEST: $nucleotides")
            for j=1:ncands
                states[j]=candidates[j][i];
            end
            if !occursin(dna"A",states)
                nucleotides[1]=DNA_A
            else
                nucleotides[1]=DNA_Gap
            end
            if !occursin(dna"T",states)
                nucleotides[2]=DNA_T
            else
                nucleotides[2]=DNA_Gap
            end
            if !occursin(dna"C",states)
                nucleotides[3]=DNA_C
            else
            nucleotides[3]=DNA_Gap
            end
            if !occursin(dna"G",states)
                nucleotides[4]=DNA_G
            else
                nucleotides[4]=DNA_Gap
            end
            #println("TEST: $states")
            push!(bases,omit_codes(nucleotides))
            #println("Test: $nucleotides")
        end
    end
    return bases
end

#given a list of base codes, return all of the base codes that do not contain any of those codes

"""
    omit_codes(bases)

Helper function for define_action_space.
"""
function omit_codes(bases::LongSequence{DNAAlphabet{4}})
    testcodes=dna"AGCTMRWSYKVHDBN"
    testcodes_copy=deepcopy(testcodes)
    indexarray=[]
    for i= 1:length(testcodes)
        count=0;
        for nucleotide in bases
            if iscompatible(nucleotide,testcodes[i])
                count+=1;
            end
        end
        if count >0
            push!(indexarray,i)
        end
    end
    check=sort(indexarray,rev=true)
    for index in check
        deleteat!(testcodes_copy,index)
    end

    return testcodes_copy
end

####################################################################
# Oligo Simulation Functions
####################################################################

"""
    isvalid(seq, sites)

Check if any of the sites appear in the sequence. Also check for palindromes.

"""
function isvalid(seq, sites)
    if is_palindrome(seq)
        return false
    end
    for site in sites
        if occursin(site, seq)
            return false
        end
    end
    return true
end

"""
    find_valid_randomer(prefix, bases, sites)

Randomly select bases to add the prefix, ensuring no site
appears in the randomer.
"""
function find_valid_randomer(prefix, bases, sites)
    n = length(bases)
    pre_n = length(prefix)
    randomer = prefix * (dna"-" ^ n)
    for i = 1:n
        candidates=shuffle(bases[i])
        for cand in candidates
            randomer[pre_n+i] = cand
            if isvalid(randomer[1:pre_n+i], sites)
                break
            else
                randomer[pre_n+i] = DNA_Gap
            end
        end
        if randomer[pre_n+i] == DNA_Gap
            break
        end
    end
    return randomer
end

############################################################
#The follwoing functions are used as part of the various reward systems. They are called by each reward function when applicable
###########################################################
"""
    measure_hitcounts(randomer::LongSequence{DNAAlphabet{4}},mRNAs)

For each mRNA, calculate the number of times the randomer appears in the sequence.
"""
function measure_hitcounts(randomer::LongSequence{DNAAlphabet{4}},mRNAs)
    mRNAnum=length(mRNAs);
    hitcounts=zeros(mRNAnum)
    for i=1:mRNAnum
        counter=0;
        indexcounter=0
        index=approxsearchindex(mRNAs[i][1:end],randomer,0);
        indexcounter+=index;
        if index > 0
            counter +=1
        end
        while index >0
            index=approxsearchindex(mRNAs[i][indexcounter+1:end],randomer,0);
            indexcounter+=index
            if index==0
                break
            else
                counter+=1;
            end
        end
        hitcounts[i]=counter
    end
    return hitcounts
end

function mRNA_uniform_coverage_function_factory(randomer::LongSequence{DNAAlphabet{4}},mRNAs,hitcounts;uniformity_weight=1,total_hits_weight=1,kwargs...)
    hitcounts+=measure_hitcounts(randomer,mRNAs)
    function mRNA_uniform_coverage(randomer,mRNAs)
        mRNAnum=length(mRNAs);
        DULQ_score=0;
        new_hitcounts=measure_hitcounts(randomer,mRNAs)
        scored_hits=hitcounts+new_hitcounts;
        if sum(scored_hits)==0
            final_score=0
        else
            sorted_hits=sort(scored_hits);
            Davg=mean(sorted_hits)
            lq=round(Integer,mRNAnum/4);
            Dlq=mean(sorted_hits[1:lq])
            DULQ_score=Dlq/Davg;
            filtered_hits=filter(x -> x > 0, scored_hits)
            final_score=length(filtered_hits)+total_hits_weight*sum(log10.(filtered_hits))+uniformity_weight*mRNAnum*DULQ_score
        end
        return final_score
    end
    function mRNA_uniform_coverage_analysis(randomer,mRNAs)
        mRNAnum=length(mRNAs);
        DULQ_score=0;
        new_hitcounts=measure_hitcounts(randomer,mRNAs)
        scored_hits=hitcounts+new_hitcounts;
        if sum(scored_hits)==0
            final_score=0
        else
            sorted_hits=sort(scored_hits);
            Davg=mean(sorted_hits)
            lq=round(Integer,mRNAnum/4);
            Dlq=mean(sorted_hits[1:lq])
            DULQ_score=Dlq/Davg;
            filtered_hits=filter(x -> x > 0, scored_hits)
            final_score=length(filtered_hits)+total_hits_weight*sum(log10.(filtered_hits))+uniformity_weight*mRNAnum*DULQ_score
        end
        return final_score,hitcounts,DULQ_score, total_hits_weight, uniformity_weight
    end
    return mRNA_uniform_coverage,mRNA_uniform_coverage_analysis,hitcounts
end

"""
    measure_hit_positions(randomer::LongSequence{DNAAlphabet{4}},mRNAs)

For every mRNA, measure the first index the randomer appears each time it appears. Store the indicies for each mRNA
"""

function measure_hit_positions(randomer::LongSequence{DNAAlphabet{4}},mRNAs)
    mRNAnum=length(mRNAs)
    hit_positions=[[] for _ =1:mRNAnum]
    for i=1:mRNAnum
        indexcounter=0
        index=approxsearchindex(mRNAs[i][1:end],randomer,0);
        indexcounter+=index;
        if index > 0
            push!(hit_positions[i],indexcounter)
        end
        while index >0
            index=approxsearchindex(mRNAs[i][indexcounter+1:end],randomer,0);
            indexcounter+=index
            if index==0
                break
            else
                push!(hit_positions[i],indexcounter)
            end
        end
    end
    return hit_positions
end
"""
    calculate_gap_distances(hit_indicies,genelen)

Calculate the distances between hit_indicies along the gene length (genelen)
"""

function calculate_gap_distances(hit_indicies,genelen)
    dists=[]
    if length(hit_indicies)>0
        hit_indicies=sort(hit_indicies)
        push!(dists,hit_indicies[1])
        if length(hit_indicies)>1
            for i=2:length(hit_indicies)
                push!(dists,(hit_indicies[i]-hit_indicies[i-1]))
            end
        end
        push!(dists,(genelen-hit_indicies[end]))
    end
    return dists
end

#hit_positions= [[] for _ in 1:mRNAnum]
function mRNA_within_gene_uniformity_function_factory(randomer::LongSequence{DNAAlphabet{4}},mRNAs,hit_positions;uniformity_weight=1,total_hits_weight=1,kwargs...)
    hit_positions=vcat.(hit_positions,measure_hit_positions(randomer,mRNAs));
    function mRNA_within_gene_uniformity(randomer,mRNAs)
        mRNAnum=length(mRNAs)
        gene_lengths=length.(mRNAs)
        new_hit_positions=measure_hit_positions(randomer,mRNAs)
        scored_hit_positions=vcat.(hit_positions,new_hit_positions)
        randomerlen=length(randomer)
        if sum(length.(scored_hit_positions))==0
            final_score=0
        else
            gap_distances=calculate_gap_distances.(scored_hit_positions,gene_lengths,randomerlen)
            uniformity_score=zeros(mRNAnum)# "gap_standard_deviations_normalized" normalized by gene length
            for i =1:mRNAnum
                if length(gap_distances[i])==0
                    uniformity_score[i]=0
                else
                    sorted_gap_distances=sort(gap_distances);
                    Davg=mean(sorted_gap_distances)
                    lq=round(Integer,length(gap_distances)/4);
                    Dlq=mean(sorted_gap_distances[1:lq])
                    uniformity_score[i]=Dlq/Davg;
                end
            end
            filtered_hits=filter(x -> length(x)>0,scored_hit_positions)
            num_hits=length.(filtered_hits)
            final_score=length(filtered_hits)+ total_hits_weight*sum(log10.(num_hits))+ uniformity_weight*sum(uniformity_score)
        end
        return final_score
    end
    function mRNA_within_gene_uniformity_analysis(randomer,mRNAs)
        mRNAnum=length(mRNAs)
        gene_lengths=length.(mRNAs)
        new_hit_positions=measure_hit_positions(randomer,mRNAs)
        scored_hit_positions=vcat.(hit_positions,new_hit_positions)
        randomerlen=length(randomer)
        uniformity_score=zeros(mRNAnum)
        if sum(length.(scored_hit_positions))==0
            final_score=0
        else
            gap_distances=calculate_gap_distances.(scored_hit_positions,gene_lengths,randomerlen)
            # "gap_standard_deviations_normalized" normalized by gene length
            for i =1:mRNAnum
                if length(gap_distances[i])==0
                    uniformity_score[i]=0
                else
                    sorted_gap_distances=sort(gap_distances);
                    Davg=mean(sorted_gap+distances)
                    lq=round(Integer,length(gap_distances)/4);
                    Dlq=mean(sorted_gap_distances[1:lq])
                    uniformity_score[i]=Dlq/Davg;
                end
            end
            filtered_hits=filter(x -> length(x)>0,scored_hit_positions)
            num_hits=length.(filtered_hits)
            final_score=length(filtered_hits)+ total_hits_weight*sum(log10.(num_hits))+ uniformity_weight*sum(uniformity_score)
        end
        return final_score,hit_positions,mean(uniformity_score),total_hits_weight,uniformity_weight
    end
    return mRNA_within_gene_uniformity,mRNA_within_gene_uniformity_analysis, hit_positions
end

"""
    calculate_DULQ_score(array)

Calculate the Distribution Uniformity, Lower Quartile score of array

DULQ = mean(25th percentile of array)/mean(array)

"""
function calculate_DULQ_score(array)
    if length(array)==0
        DULQ_score=0
    else
        sorted_array=sort(array)
        grand_mean=mean(sorted_array)
        lq_idx=max(round(Integer,length(sorted_array)/4),1);
        lq_mean=mean(sorted_array[1:lq_idx])
        DULQ_score=lq_mean/grand_mean;
    end
    return DULQ_score
end
"""
    mRNA_all_rewards_function_factory(randomer::LongSequence{DNAAlphabet{4}},mRNAs,hitcounts,hit_positions;genes_hit_weight=1,intrauniformity_weight=0,total_hits_weight=0,interuniformity_weight=0,kwargs...)

Create reward and analysis functions to be used during rollout simulation.

# Arguments
- `randomer`: Oligo to be added to the pool. The previously chosen cadidate.
- `mRNAs`: Array of mRNA sequences to be hit by NSR primers
- `hitcounts`: For each mRNA, the number of hits from previous oligos.
- `hit_positions`: For each mRNA, the index of each hit from previous oligos
- `genes_hit_weight=1`: Weight for hitting every gene at least once
- `total_hits_weight=0`: Weight for maximizing total hits
- `interuniformity_weight=0`: Weight for distributing equal numbers of hits to each gene
- `intrauniformity_weight=0`: Weight for placing hits uniformly along each gene. If 0, we do not track hit positions to reduce unnecessary computation.
"""
function mRNA_all_rewards_function_factory(randomer::LongSequence{DNAAlphabet{4}},mRNAs,hitcounts,hit_positions;genes_hit_weight=1,intrauniformity_weight=0,total_hits_weight=0,interuniformity_weight=0,kwargs...)
    if intrauniformity_weight==0  # we will only track the hit positions if the intrauniformity score matters. This will save unnecessary computation
        hitcounts+=measure_hitcounts(randomer,mRNAs)

        function mRNA_all_rewards_A(randomer,mRNAs)
            mRNAnum=length(mRNAs);
            interuniformity_score=0;
            new_hitcounts=measure_hitcounts(randomer,mRNAs)
            scored_hits=hitcounts+new_hitcounts;
            if sum(scored_hits)==0
                final_score=0
            else
                filtered_hits=filter(x -> x > 0, scored_hits)
                interuniformity_score=calculate_DULQ_score(filtered_hits)
                final_score=genes_hit_weight*length(filtered_hits)+total_hits_weight*sum(filtered_hits)+interuniformity_weight*mRNAnum*interuniformity_score
            end
            return final_score
        end
        function mRNA_all_rewards_analysis_A(randomer,mRNAs)
            mRNAnum=length(mRNAs);
            interuniformity_score=0;
            new_hitcounts=measure_hitcounts(randomer,mRNAs)
            scored_hits=hitcounts+new_hitcounts;
            genes_hit=0
            total_hits=0
            interuniformity_score=0
            intrauniformity_score=0
            if sum(scored_hits)==0
                final_score=0
                genes_hit=0
                total_hits=0
                interuniformity_score=0
                intrauniformity_score=0;
            else
                filtered_hits=filter(x -> x > 0, scored_hits)
                interuniformity_score=calculate_DULQ_score(filtered_hits)
                genes_hit=length(filtered_hits)
                total_hits=sum(filtered_hits)
                intrauniformity_score= -1
                final_score=genes_hit_weight*length(filtered_hits)+total_hits_weight*sum(filtered_hits)+interuniformity_weight*mRNAnum*interuniformity_score
            end
            return final_score,hitcounts,hit_positions,genes_hit,total_hits,interuniformity_score,mean(intrauniformity_score),genes_hit_weight,total_hits_weight,interuniformity_weight,intrauniformity_weight
        end
        mRNA_all_rewards=mRNA_all_rewards_A
        mRNA_all_rewards_analysis=mRNA_all_rewards_analysis_A

    else ## intrauniformity_weight >0, keep track of hit positions
        hit_positions=vcat.(hit_positions,measure_hit_positions(randomer,mRNAs));
        hitcounts=length.(hit_positions)
        function mRNA_all_rewards_B(randomer,mRNAs)
            mRNAnum=length(mRNAs)
            gene_lengths=length.(mRNAs)
            new_hit_positions=measure_hit_positions(randomer,mRNAs)
            scored_hit_positions=vcat.(hit_positions,new_hit_positions)
            randomerlen=length(randomer)
            if sum(length.(scored_hit_positions))==0
                final_score=0
            else
                gap_distances=calculate_gap_distances.(scored_hit_positions,gene_lengths,randomerlen)
                gap_distances=filter(x -> length(x)>0,gap_distances)
                intrauniformity_scores=calculate_DULQ_score.(gap_distances)
                filtered_hit_positions=filter(x -> length(x)>0,scored_hit_positions)
                filtered_hitcounts=(length.(filtered_hit_positions));
                interuniformity_score=calculate_DULQ_score(filtered_hitcounts)
                filtered_hits=length.(filtered_hit_positions)
                final_score=genes_hit_weight*length(filtered_hits)+ total_hits_weight*sum(filtered_hits)+ interuniformity_weight*mRNAnum*interuniformity_score+intrauniformity_weight*mRNAnum*mean(intrauniformity_scores)
            end
            return final_score
        end
        function mRNA_all_rewards_analysis_B(randomer,mRNAs)
            mRNAnum=length(mRNAs)
            gene_lengths=length.(mRNAs)
            new_hit_positions=measure_hit_positions(randomer,mRNAs)
            scored_hit_positions=vcat.(hit_positions,new_hit_positions)
            randomerlen=length(randomer)
            genes_hit=0
            total_hits=0
            interuniformity_score=0
            intrauniformity_scores=0
            if sum(length.(scored_hit_positions))==0
                final_score=0
                genes_hit=0
                total_hits=0
                interuniformity_score=0
                intrauniformity_scores=0
            else
                gap_distances=calculate_gap_distances.(scored_hit_positions,gene_lengths,randomerlen)
                gap_distances=filter(x -> length(x)>0,gap_distances)
                intrauniformity_scores=calculate_DULQ_score.(gap_distances)
                filtered_hit_positions=filter(x -> length(x)>0,scored_hit_positions)
                filtered_hitcounts=(length.(filtered_hit_positions));
                interuniformity_score=calculate_DULQ_score(filtered_hitcounts)
                filtered_hits=length.(filtered_hit_positions)
                genes_hit=length(filtered_hits)
                total_hits=sum(filtered_hits)
                final_score=genes_hit_weight*length(filtered_hits)+ total_hits_weight*sum(filtered_hits)+ interuniformity_weight*mRNAnum*interuniformity_score+intrauniformity_weight*mRNAnum*mean(intrauniformity_scores)
            end
            return final_score,hitcounts,hit_positions,genes_hit,total_hits,interuniformity_score,mean(intrauniformity_scores),genes_hit_weight,total_hits_weight,interuniformity_weight,intrauniformity_weight
        end
        mRNA_all_rewards=mRNA_all_rewards_B
        mRNA_all_rewards_analysis=mRNA_all_rewards_analysis_B
    end
    return mRNA_all_rewards,mRNA_all_rewards_analysis,hitcounts,hit_positions
end
#################################################
# Simulation function. Call Reward functions and simulate a valid randomer
##################################################

"""
    simulate(prefix,bases,sites,mRNAs,reward_function; nsims=1000, horizon=length(bases), kwargs...)

Perfrom rollout simulations and score those simulations

#Arguments
- prefix: The established portion of the randomer along with the selected base to perform rollout on
- bases: The remaining availabe base codes at each position to be chosen at random during simulation
- mRNAs: mRNA sequences used for scoring
- reward function: scoring funciton to be used in evaluating random simulations
- nsims=1000: number of simulations per rollout
- horizon: how far ahead to look with simulations
"""
function simulate(prefix,bases,sites,mRNAs,reward_function; nsims=1000, horizon=length(bases), kwargs...)
    rewards=zeros(nsims)
    true_horizon = min(horizon, length(bases))
    n=length(prefix)+true_horizon;
    sim_randomer_list=[dna"-"^n for _ in 1:nsims];

    Threads.@threads for i = 1:nsims   #This is where multithreading is incorporated
        sim_randomer_list[i]=find_valid_randomer(prefix,bases[1:true_horizon],sites);
        rewards[i]= reward_function(sim_randomer_list[i],mRNAs)
    end
    max_sim_score,max_index=findmax(rewards);
    max_sim_randomer=sim_randomer_list[max_index];
    return mean(rewards),max_sim_randomer,max_sim_score

end

########################################
# Rollout Function
#######################################

"""
    oligo_rollout(bases,sites, mRNA; reward_function=mRNA_uniform_coverage, kwargs...)

Select a randomer using the rollout algorithm

#Arguments
- bases: Available base codes for each position of the desired randomer
- sites: Array of sequences that aren't allowed in the randomer. Each site must be be the same length as the desired randomer
- mRNA: Array of mRNA sequences the randomer is trying to capture
-reward_function=mRNA_uniform_coverage: Reward funciton to be used by rollout algorithm.
"""
function oligo_rollout(bases,sites, mRNA; reward_function=mRNA_uniform_coverage, kwargs...)
    n=length(bases)
    randomer= dna"-" ^ n
    # Find the longest blocking site and use this as the horizon. (they should all be equal)
    max_len = map(length, sites) |> maximum
    candidates=valid_candidates(sites);
    best_sim_randomer=dna"-"^n
    best_sim_score=0;
    for i= 1:n
        best_base= DNA_Gap
        best_score=0;
        start = max(i - max_len, 1) # only need to look back max_len when checking the oligos
        bases=define_action_space(candidates[findall(j -> iscompatible(j[start:i-1],randomer[start:i-1]),candidates)])
        if length(bases[i])==1
            best_base=bases[i][1];
        else
            for base in bases[i]
                if ~isvalid(randomer[start:i-1] * base, sites)
                    # adding this bases creates a restriction site; skip it
                    continue
                end
                score,max_sim_randomer,max_sim_score= simulate(randomer[start:i-1] * base, bases[i+1:end], sites, mRNA,reward_function;kwargs...)
                # Update use this base if:
                #   1. the mean randomer mRNA coverage is higher than the previous best, OR
                #   2. the randomer mRNA coverage is tied with the previous best, but the new base has
                #      higher individual degeneracy
                if score > best_score || (score == best_score && degeneracy(base) > degeneracy(best_base))
                    best_score = score
                    best_base = base
                end
                ## Update the overall best simulation result if
                #   1. the best globally simulated randomer mRNA score is better than the previous best
                #   2. the best globally simulated randomer mRNA score is tied but the randomer has higher degeneracy
                if max_sim_score >best_sim_score || (max_sim_score==best_sim_score && degeneracy(max_sim_randomer)>degeneracy(best_sim_randomer))
                    best_sim_score=max_sim_score;
                    best_sim_randomer=max_sim_randomer
                end
            end
        end
        if best_base == DNA_Gap
            # The sequence has terminated; there are no bases that can be added without making
            # a restriction site.
            break
        else
            randomer[i] = best_base

        end
        #println("randomer: $randomer")
    end
    final_randomer_score=reward_function(randomer,mRNA)
    if best_sim_score >final_randomer_score
        randomer=best_sim_randomer
        println("Triggered Best Simulated Randomer")
    end
    return randomer
end

####################################
#Functions for Reading Genome Data CSV Files Imported from R
######################################

## Read in the list of sequences you want to exclude in the designed primers.
# Function finds every unique hexamer in the list of sequences and stores it as an array
#. For example, use list of rRNA and tRNA sequences. This funciton takes care of orientation. We end up using the reverse complement since we are finding primers that won't hybridize
## Example
#blocking_sites=read_blocking_sites(CSV.read("./SMU_UA159_rRNA_tRNA.csv"))

"""
    read_blocking_sites(data; randomerlen=6)

Read in the list of unwanted sequences and store every unique randomerlen-mer reverse complement sequence.
"""
function read_blocking_sites(data; randomerlen=6)
    blocking_sites=[];
    for i =1:nrow(data)
        seq=reverse_complement(LongDNASeq(string(data[i,:Sequence])));
        for j=1:length(seq)-(randomerlen-1)
            if isvalid(seq[j:j+(randomerlen-1)],blocking_sites)
                push!(blocking_sites,seq[j:j+(randomerlen-1)])
            end
        end
    end
    return blocking_sites
end
# RNAs=read_rRNA_tRNA_list(CSV.read("./SMU_UA159_rRNA_tRNA.csv"))
#This function Does not hash up the rRNAs and tRNAs like read_blocking_sites does. That is the only difference

"""
    read_rRNA_tRNA_list(data)

Read in the list of unwanted sequences and store every unique reverse_complement sequence
"""
function read_rRNA_tRNA_list(data)
    RNAs=[]
    for i=1:nrow(data)
        seq=reverse_complement(LongDNASeq(string(data[i,:Sequence])));
        if isvalid(seq,RNAs)
            push!(RNAs,seq)
        end
    end
    return RNAs
end
## Read in the list of sequences you want the designed primers to hybridize to
# Function takes care of the orientation, we want to end up using the reverse complmement.
## Example
# mRNAs=read_mRNA_list(CSV.read("./SMU_UA159_Genes.csv"))

"""
    read_mRNA_list(data)

Read in the list of target sequences and store every unique reverse_complement sequence
"""
function read_mRNA_list(data)
    mRNAs=[]
    for i=1:nrow(data)
        seq=reverse_complement(LongDNASeq(string(data[i,:mRNA_seq])));
        if isvalid(seq,mRNAs)
            push!(mRNAs,seq)
        end
    end
    return mRNAs
end

#Finds the number of bases as well as the GC content of the given file
# Example
# bases,GC_content=calculate_transcriptome_stats(read_mRNA_list(CSV.read("./SMU_UA159_Genes.csv")))
"""
    calculate_transcriptome_stats(sequences)

Find the number of bases and total GC content of sequences
"""
function calculate_transcriptome_stats(sequences)
    bases=0;  #number of bases in the file
    GC_count=0 # Count of G's and C's in file
    for i=1:length(sequences)
        bases += length(sequences[i]);
        comp=composition(sequences[i]);
        GC_count += comp[DNA_G] +comp[DNA_C];
    end
    GC_content= GC_count/bases*100;
    return bases,GC_content
end

#################################################
#### Extra Analysis Tools
#################################################

## Finds the edit distance distribution for a given designed primer to all of the blocking sites
function base_error_dist(randomer,sites)
    error_dist=[]
    for i in 0:length(randomer)
        count=0
        found_sites=[]
        for j in 1:length(sites)
            valid=approxsearchindex(randomer,sites[j],i)
            if valid == 1
                count += valid
                push!(found_sites,j)
            end
        end
        deleteat!(sites,found_sites)
        push!(error_dist,count)
    end

    return error_dist
end

"""
    ID_genes_missing(data,genes_missing)

Find the Gene_ID within data given the indicies of genes_missing
"""
function ID_genes_missing(data,genes_missing)
    missing_ID=[];
    for i in genes_missing
        push!(missing_ID,data[i,:New_ID])
    end
    return missing_ID
end

## Calculates the distribution number of randomers that hit each gene at least once. Not very useful for most cases
function coverage_distribution(randomers, mRNA)
    coverage_dist=zeros(length(randomers)+1)
     for seq in mRNA
         count=0;
         for randomer in randomers
             if occursin(randomer, seq)
                 count+=1;
             end
         end
         coverage_dist[count+1]+=1;

     end
     return coverage_dist
end

# Calcualtes the number of times each gene is hit by all randomers in the pool. Useful for comparing the "between gene uniformity".
function gene_coverage_counts(randomers,mRNAs)
    gene_hits=zeros(length(mRNAs))
    randomerlen=length(randomers[1]);
    for i in 1:length(mRNAs)
        counter=0;
        for j in 1:length(mRNAs[i])-randomerlen
            for k in 1:length(randomers)
                if occursin(randomers[k],mRNAs[i][j:j+randomerlen-1])
                    counter+=1
                end
            end
        end
        gene_hits[i]=counter
    end
    data=DataFrame(counts=gene_hits)
    return data

end

 ########################################################################
 # MAIN RUNNING FUNCTION
 ########################################################################

 #Quick Test#
 # data=NSR_RL(;pool_size=3,nsims=50)
 #CSV.write("NSR_RL_S_mutans_4_12_21",data)

 #Input the Species name, the length of primer, and the size of the pool.

 """
    NSR_RL(species,mRNAfile,rRNAfile,randomerlen,pool_size,nucleotides = dna"AGCTMRWSYKVHDBN";kwargs...)

Select randomers to selectively prime reverse transcription

#Arguments
- `mRNAfile`: The CSV file from Genome_Info_Mining that contains the mRNA sequences
- `rRNAtRNAfile`: The CSV file from Geneom_Info_Mining that contains the rRNA/tRNA sequences
- `randomerlen`: The length of the primers to be designed
- `pool_size`: The number of primers to be made
- `nucleotides=dna"AGCTMRWSYKVHDBN"`: The base set of available codes to be used for designing primers.

#Optional Keyword Arguments
- `nsims`: The number of rollout simulations per action
- `genes_hit_weight`: Weight given to hitting each gene at least once in the reward function
- `total_hits_weight`: Weight given to maximizing the total number of hits
- `interuniformity_weight`: Weight given to placing hits equally across all genes
- `intrauniformity_weight`: Weight given to placing hits equally within each gene
 """
function NSR_RL(species,mRNAfile,rRNAfile,randomerlen,pool_size,nucleotides = dna"AGCTMRWSYKVHDBN";kwargs...)

    #Intialization
    randomers=[]
    mRNA_DF=CSV.read(mRNAfile,DataFrame);
    rRNA_DF=CSV.read(rRNAfile,DataFrame);
    mRNAs=read_mRNA_list(mRNA_DF);
    rRNAs=read_rRNA_tRNA_list(rRNA_DF);
    numgenes=length(mRNAs)
    blocking_sites=read_blocking_sites(rRNA_DF,randomerlen=randomerlen);
    mRNAbp,mRNA_GC=calculate_transcriptome_stats(mRNAs);
    rRNAbp,rRNA_GC=calculate_transcriptome_stats(rRNAs);
    transcriptome_size=mRNAbp+rRNAbp;
    all_ns = [nucleotides for _ in 1:randomerlen];
    times=[]
    num_blocked_sites=zeros(pool_size)
    println("Initialization Successful")
    ##Initialize the reward-associated data structure and reward function here. ex hitcounts=zeros(numgenes), hit_positions=[[] for _ in 1:numgenes]
    hitcounts=zeros(numgenes)
    hit_positions=[[] for _ in 1:numgenes]
    #Call the appropriate reward function factory
    reward_function,analysis_function,hitcounts,hit_positions=mRNA_all_rewards_function_factory(dna"-"^randomerlen,mRNAs,hitcounts,hit_positions;kwargs...)
    ##
    for i = 1:pool_size
        num_blocked_sites[i]=length(blocking_sites)
        randomerstats= @timed oligo_rollout(all_ns,blocking_sites,mRNAs; reward_function=reward_function,kwargs...); #create new randomer and add to pool
        push!(randomers,randomerstats.value)
        push!(times,randomerstats.time)
        push!(blocking_sites,LongDNASeq(randomers[i]));
        push!(blocking_sites,LongDNASeq(reverse_complement(randomers[i]))); ## avoid adding randomers that are reverse complements of each other
        blocking_sites=unique(decompress_pool(blocking_sites));
        reward_function,analysis_function,hitcounts,hit_positions=mRNA_all_rewards_function_factory(randomers[i],mRNAs,hitcounts,hit_positions; kwargs...) #officially add hit positions from the new randomer to the reward function
        #println("Iteration $i")
    end
    final_score,hitcounts,hit_positions,genes_hit,total_hits,interuniformity_score,intrauniformity_score,genes_hit_weight,total_hits_weight,interuniformity_weight,intrauniformity_weight =analysis_function(dna"-"^randomerlen,mRNAs)
    genes_missing=findall(x -> x==0,hitcounts)
    missing_IDs=join(ID_genes_missing(mRNA_DF,genes_missing),",")
    totaldeg=pooled_degeneracy(randomers)
    total_time=sum(times)
    string_time=join(times,",")
    string_num_blocked_sites=join(num_blocked_sites,",")
    string_randomers=join(convert.(String,randomers),",")
    data=DataFrame(poolsize=pool_size, primer_length=randomerlen, genes_hit_weight=genes_hit_weight, total_hits_weight=total_hits_weight,interuniformity_weight=interuniformity_weight,intrauniformity_weight=intrauniformity_weight,species=species,num_genes=numgenes,
        transcriptome_size=transcriptome_size,mRNA_GC=mRNA_GC,blocked_GC=rRNA_GC,randomers=string_randomers,degeneracy=totaldeg,num_genes_hit=genes_hit,
        genes_missing=missing_IDs,total_hits=total_hits,interuniformity_score=interuniformity_score,intrauniformity_score=intrauniformity_score,individual_times=string_time,time=total_time,num_blocked_sites=string_num_blocked_sites,final_score=final_score)
    return data
end

function parse_commandline()
    s = ArgParseSettings()

    @add_arg_table s begin
        "--accession","-a"
            help = "Geneome Accesion number. Separate multiple accession numbers by comma wiht no spaces"
        "--species", "-s"
            help = "species name"
            default = "AGCTMRWSYKVHDBN"
            required = true
        "--genefile"
            help = "gene file name"
        "--rRNAtRNAfile"
            help = "rRNA/tRNA file name"
        "--nsims","-n"
            help = "Number of simulations per rollout"
            arg_type= Int
            default=100
        "--genes_hit_weight"
            help = "Weight given to the reward for hitting every gene at least once"
            default=1
        "--total_hits_weight"
            help = "Weight given to the reward for maximizing total hits"
            default=1
        "--intrauniformity_weight"
            help = "Weight given to placing hits uniformly between genes"
            default=1
        "--interuniformity_weight"
            help = "Weight given to placing hits uniformly within genes"
            default=1
        "--bases","-b"
            help = "Available base codes for the primers to contain. default=AGCTMRWSYKVHDBN"
            default= "AGCTMRWSYKVHDBN"
        "--length_primers","-l"
            help = "Primer Length. Default=6"
            default=6
        "--poolsize","-p"
            help= "Size of primer pool"
            arg_type=Int
            required=true
        "--outfile","-o"
            help="Output file names"

    end

    return parse_args(s)
end

function main()
    args = parse_commandline()
    println("Parsed args:")
    for (arg,val) in args
        println("  $arg  =>  $val")
    end
    accession=args["accession"]
    species=args["species"]
    genefile=args["genefile"]
    rRNAtRNAfile=args["rRNAtRNAfile"]
    nsims=args["nsims"]
    genes_hit_weight=args["genes_hit_weight"]
    total_hits_weight=args["total_hits_weight"]
    intrauniformity_weight=args["intrauniformity_weight"]
    interuniformity_weight=args["interuniformity_weight"]
    bases=args["bases"]
    bases=LongDNASeq(bases)
    length_primers=args["length_primers"]
    poolsize=args["poolsize"]
    if accession===nothing && (genefile===nothing || rRNAtRNAfile===nothing)
        error("Please Provide either an accession number(s) or a gene file AND a rRNA/tRNA file")
    end
    if accession===nothing
        println("Using Provided Files")
    else
        println("Using Accession number")
        @rput accession
        @rput species
        R"""
        library(nsrgenomes)
        nsrgenomes(species,accession)
        """
        genefile=string(species,"_Genes.csv")
        rRNAtRNAfile=string(species,"_rRNA_tRNA.csv")
    end
    DF=NSR_RL(species,genefile,rRNAtRNAfile,length_primers,poolsize,bases;nsims=nsims,genes_hit_weight=genes_hit_weight,total_hits_weight=total_hits_weight,intrauniformity_weight=intrauniformity_weight,interuniformity_weight=interuniformity_weight)
    str_randomers=DF[1,12]
    randomers=split(str_randomers,",")
    if outfile===nothing
        show(stdout,"text/plain",randomers)
        println("\n")
    else
        f=open(outfile,"w")
        for seq in randomers
            println(f,seq)
        end
        close(f)
    end
end

main()


text:
using YAML

dbfolder = Pkg.dir("GlassCatalog","assets","refractiveindex.info-database","database")
dbfile = joinpath(dbfolder, "library.yml")
library = YAML.load_file(dbfile)

using DataArrays, DataFrames

shelves = Dict{UTF8String,UTF8String}()
books = Dict{UTF8String,UTF8String}()
pages = Dict{UTF8String,UTF8String}()

libraryData = DataFrame(shelf=UTF8String[], subset=UTF8String[], book=UTF8String[], page=UTF8String[], path=UTF8String[])

for shelf in library
  shelfName::UTF8String = shelf["SHELF"]
  shelfDesc::UTF8String = shelf["name"]
  setindex!(shelves, shelfName, shelfDesc)
  subset::UTF8String = ""
  for book in shelf["content"]
    if haskey(book,"DIVIDER")
      subset = book["DIVIDER"]
    else
      bookName::UTF8String = book["BOOK"];
      bookDesc::UTF8String = book["name"]
      setindex!(books, bookName, bookDesc)
      for page in book["content"]
        if haskey(page, "PAGE")
          pageName::UTF8String = page["PAGE"]
          pageDesc::UTF8String = page["name"]
          pagePath::UTF8String = page["path"]
          setindex!(pages, pageName, pageDesc)
          pageFile = joinpath(dbfolder, pagePath)
          pageData = YAML.load_file(pageFile)
          df = DataFrame(shelf=[shelfName], subset=[subset], book=[bookName], page=[pageName], path=[pagePath], data=[pageData])
          append!(libraryData, df)
        end
      end
    end
  end
end


text:
module Newton
using LinearAlgebra
using RecursiveFactorization
using DiffResults
using ForwardDiff
using StaticArrays

struct NewtonCache{Tres,Tcfg}
    result::Tres
    config::Tcfg
    lupivot::Vector{Int}
end

"""
    function NewtonCache(x::AbstractVector, rf!)

Create the cache used by the `newtonsolve!` and `linsolve!`
to find `x` such that `rf!(r,x)` yields `r=0`.
"""
function NewtonCache(x::AbstractVector, rf!)
    result = DiffResults.JacobianResult(x)
    cfg = ForwardDiff.JacobianConfig(rf!, x, result.value, ForwardDiff.Chunk(length(x)))
    lupivot = Vector{Int}(undef, length(x))
    return NewtonCache(result, cfg, lupivot)
end

"""
    get_drdx(cache::NewtonCache) = DiffResults.jacobian(cache.result)

"""
get_drdx(cache::NewtonCache) = DiffResults.jacobian(cache.result)

"""
    linsolve!(K::AbstractMatrix, b::AbstractVector, cache::NewtonCache)

Solves the linear equation system `Kx=b`, mutating both `K` and `b`.
`b` is mutated to the solution `x`
"""
function linsolve!(K::AbstractMatrix, b::AbstractVector, cache::NewtonCache)
    LU = RecursiveFactorization.lu!(K, cache.lupivot)
    ldiv!(LU, b)
    return b
end

"""
    newtonsolve!(x::AbstractVector, drdx::AbstractMatrix, rf!, cache::ResidualCache; tol=1.e-6, maxiter=100)

Solve the nonlinear equation system r(x)=0 using the newton-raphson method. Returns `true` if converged and `false` otherwise.

# args
- `x`: Vector of unknowns. Provide as initial guess, mutated to solution.
- `drdx`: Jacobian matrix. Only provided as preallocation. Can be aliased to `DiffResults.jacobian(cache.result)`
- `rf!`: Residual function. Signature `rf!(r, x)` and mutating the residual `r`
- `cache`: Optional cache that can be preallocated by calling `ResidualCache(x, rf!)`

# kwargs
- `tol=1.e-6`: Tolerance on `norm(r)`
- `maxiter=100`: Maximum number of iterations before no convergence

"""
function newtonsolve!(x::AbstractVector, drdx::AbstractMatrix, rf!, cache::NewtonCache = NewtonCache(x,rf!); tol=1.e-6, maxiter=100)
    diffresult = cache.result
    cfg = cache.config
    for i = 1:maxiter
        # Disable checktag using Val{false}(). solve_residual should never be differentiated using dual numbers!
        # This is required when using a different (but equivalent) anynomus function for caching than for running.
        ForwardDiff.jacobian!(diffresult, rf!, diffresult.value, x, cfg, Val{false}())
        err = norm(DiffResults.value(diffresult))
        i == 1 && @assert !(typeof(err) <: ForwardDiff.Dual)    # Check that we don't try to differentiate
        if err < tol
            drdx .= DiffResults.jacobian(diffresult)
            return true
        end
        linsolve!(DiffResults.jacobian(diffresult), DiffResults.value(diffresult), cache)
        x .-= DiffResults.value(diffresult)
    end
    # No convergence
    return false
end

"""
    linsolve(drdx::SMatrix{dim,dim}, r::SVector{dim}) where{dim}

Solves the linear equation system `drdx*x=r` without mutating and returns the solution x
"""
@inline linsolve(drdx::SMatrix{dim,dim}, r::SVector{dim}) where{dim} = drdx\r

"""
    newtonsolve(x0::SVector, rf; tol=1.e-6, maxiter=100)

Solve the nonlinear equation system `r(x)=0` using the newton-raphson method.
Returns type: `(converged, x, drdx)`, SVector, SMatrix)` where
- `converged::Bool` is `true` if converged and `false` otherwise
- `x::SVector` is the solution vector such that `r(x)=0`
- `drdx::SMatrix` is the jacobian at `x`

# args
- `x0`: Vector of with initial guess for unknowns.
- `rf`: Residual function. Signature `r=rf(x::SVector{dim})::SVector{dim}`

# kwargs
- `tol=1.e-6`: Tolerance on `norm(r)`
- `maxiter=100`: Maximum number of iterations before no convergence

"""
function newtonsolve(x::SVector{dim}, rf; tol=1.e-6, maxiter=100) where{dim}
    for _ = 1:maxiter
        r = rf(x)
        err = norm(r)
        drdx = ForwardDiff.jacobian(rf, x)
        if err < tol
            return true, x, drdx
        end
        x -= linsolve(drdx, r)
    end
    return false, zero(SVector{dim})*NaN, zero(SMatrix{dim,dim})*NaN
end

export newtonsolve!, newtonsolve
export NewtonCache
export get_drdx

end


text:
#########################################################
### World Geodetic Coordinate System of 1984 (WGS 84) ###
### Standardized coordinate system for Earth          ###
### Global ellipsoidal reference surface              ###
#########################################################

const WGS84  = OpenStreetMapX.Ellipsoid(a = 6378137.0, f_inv = 298.257223563)
const OSGB36 = OpenStreetMapX.Ellipsoid(a = 6377563.396, b = 6356256.909)
const NAD27  = OpenStreetMapX.Ellipsoid(a = 6378206.4,   b = 6356583.8)

#########################
### Point Translators ###
#########################

getX(lla::OpenStreetMapX.LLA) = lla.lon
getY(lla::OpenStreetMapX.LLA) = lla.lat
getZ(lla::OpenStreetMapX.LLA) = lla.alt

getX(enu::OpenStreetMapX.ENU) = enu.east
getY(enu::OpenStreetMapX.ENU) = enu.north
getZ(enu::OpenStreetMapX.ENU) = enu.up

################
### Distance ###
################

# Point translators
distance(a::OpenStreetMapX.ENU, b::OpenStreetMapX.ENU) = OpenStreetMapX.distance(a.east, a.north, a.up,
                                    b.east, b.north, b.up)

distance(a::OpenStreetMapX.ECEF, b::OpenStreetMapX.ECEF) = OpenStreetMapX.distance(a.x, a.y, a.z,
                                      b.x, b.y, b.z)

function distance(x1, y1, z1, x2, y2, z2)
    return sqrt((x2-x1)^2 + (y2-y1)^2 + (z2-z1)^2)
end


text:
const NOTEBOOKS_DIR = joinpath(pkgdir(JModels), "docs", "src", "notebooks")

"""Run all Pluto notebooks (".jl" files) in `NOTEBOOKS_DIR` and write outputs to files."""
function build()
    println("Building notebooks")
    # Evaluate notebooks in the same process to avoid having to recompile from scratch each time.
    # This is similar to how Documenter and Franklin evaluate code.
    # Note that things like method overrides and other global changes may leak between notebooks!
    use_distributed = false
    output_format = documenter_output
    bopts = BuildOptions(NOTEBOOKS_DIR; use_distributed, output_format)
    build_notebooks(bopts)
    return nothing
end

"Return Markdown file links which can be passed to Documenter.jl."
function markdown_files()
    md_files = map(notebooks) do tutorial
        file = lowercase(replace(tutorial, " " => '_'))
        return joinpath("notebooks", "$file.md")
    end
    return md_files
end


text:
struct NRWork
	D::Array{Float64, 2}
	tmp_m::Vector{Float64}
	tmp_m2::Vector{Float64}
	dc::Vector{Float64}
end

NRWork(m::Int) = NRWork(Array{Float64}(undef, m, m), [Array{Float64}(undef, m) for i in 1:3]...)

mutable struct NR
	U::Array{Float64, 2}
	Σ::Array{Float64, 1}
	Vt::Array{Float64, 2}
	tol::Float64
	maxiter::Int
	work::NRWork
	ineq::Bool
	idata::InequalityData
end

struct ProjPenaltyWork
	J::Array{Float64, 2}
	tmp_m::Vector{Float64}
	r::Vector{Float64}
	p::Vector{Float64}
	z::Vector{Float64}
	dx::Vector{Float64}
	g::Vector{Float64}
	cvalaug::Vector{Float64}
end

ProjPenaltyWork(m::Int, n::Int, m_ineq::Int, n_ineq::Int) = ProjPenaltyWork(Array{Float64}(undef, m, n),
	Array{Float64}(undef, m_ineq), [Array{Float64}(undef, n_ineq) for i in 1:5]..., Array{Float64}(undef, m_ineq))

mutable struct ProjPenalty{F}
	jac!::F
	U::Array{Float64, 2}
	Σ::Array{Float64, 1}
	Vt::Array{Float64, 2}
	rank::Int
	μ0::Float64
	tol::Float64
	maxiter::Int
	maxiter_pcg::Int
	work::ProjPenaltyWork
	ineq::Bool
	idecomp::InequalityDecomp
	idata::InequalityData
end

struct Euclidean
end

struct YRetract
	idata::InequalityData
end

# -------------------------- Retraction functions -------------------------------------

function retract!(cval, xnew, c!, xtilde, x, method::Euclidean)
	xnew .= xtilde

	return 0, 0, 0
end

function retract!(cval, xnew, c!, xtilde, x, method::YRetract)
	xnew .= xtilde
	y_retract!(xnew, x, method.idata)

	return 0, 0, 0
end

function retract!(cval, xnew, c!, xtilde, x, method::NR)
	#= performs Newton-Raphson retraction for c(xtilde + U d) using the Jacobian V Σ'
	which is evaluated at x

	termination criterion is tolerance in inf norm of c!

	INPUT
	cval - contraint value vector
	xnew - vector in which to store the result
	c! - constraint function of the form c!(cval, x) where y is overwritten
	xtilde - current value of x at step in tangent space
	x - previous iterate
	U, Σ, Vt - SVD decomposition
	tol - function tolerance for termination
	maxiter - maximum number of allowable iterations
	work - NRWork struct for storage

	OUTPUT
	flag - 0 = success, 1 = maxiter reached
	i - total number of iterations taken

	OVERWRITTEN
	cval, xnew

	=#

	# unpack
	U = method.U
	Σ = method.Σ
	Vt = method.Vt
	tol = method.tol
	maxiter = method.maxiter
	work = method.work

	# extract quantities from NRWork
	D = work.D    			# m x m array to store inverse of Jacobian of c(xtilde + Ud)
	tmp_m = work.tmp_m
	tmp_m2 = work.tmp_m2
	dc = work.dc

	m = length(Σ)
	xnew .= xtilde

	if method.ineq
		y_retract!(xnew, x, method.idata)
		c!(cval, view(xnew, 1:length(xnew)÷2))
	else
		c!(cval, xnew)
	end

	# calculate inverse Jacobian Σ^(-1) V'
	for j in 1:m
		@fastmath @inbounds @simd for k in 1:m
			D[k,j] = Vt[k,j] / Σ[k]
		end
	end

	i = 0
	while i < maxiter
		# check if tolerance met
		if norm(cval, Inf) < tol
			break
		end

		# take Newton-Raphson step - tmp_m stores step
		gemv!('N', -1.0, D, cval, 0.0, tmp_m)	# d = - (V Σ') \ cval = - Σ^(-1) V' cval
		gemv!('N', 1.0, U, tmp_m, 1.0, xnew)	# xnew = xnew + U d

		# update function values
		if method.ineq
			y_retract!(xnew, x, method.idata)
			c!(tmp_m2, view(xnew, 1:length(xnew)÷2))
		else
			c!(tmp_m2, xnew)
		end

		# Broyden update
		dc .= tmp_m2 .- cval
		cval .= tmp_m2

		# Good Broyden
		gemv!('T', 1.0, D, tmp_m, 0.0, tmp_m2) # D' dx
		gemv!('N', -1.0, D, dc, 1.0, tmp_m)		# tmp_m = Δx - D*Δc

	    alpha = 1 / dot(tmp_m2, dc)
	    ger!(alpha, tmp_m, tmp_m2, D)

	    # # Bad Broyden
	    # gemv!('N', -1.0, D, dc, 1.0, tmp_m)		# tmp_m = Δx - D*Δc

	    # alpha = 1 / dot(dc, dc)
	    # ger!(alpha, tmp_m, dc, D)

		i += 1
	end

	flag = 0
	if i == maxiter
		flag = 1
	end

	return flag, i, 0
end

function pcg!(μ, J, M!, x, r, p, z, tmp_m, tol, maxiter)
	#= performs preconditioned conjugate gradient to solve Ax = b with
		A = J' J + μI
		preconditioned with M!

		INPUT
		μ - factor for projection
		J - the Jacobian of the constraints
		M! - function for preconditioning of the form M!(z, r)
		x - initial guess
		r - residual, should equal b (the right-hand side) using an initial guess of x=0
		p, z - work vectors of size n for PCG
		tmp_m - work vector of size m
		tol - tolerance for convergence
		maxiter - maximum number of iterations allowed

		OUTPUT
		x, r, p, z, tmp_m overwritten
		x - overwritten with solution
		r - overwritten with residual
		flag - 0 = success, 1 = maxiter reached
	=#

	norm_res = Inf
	ρ = 1.0               	# set ρ and p initialy to 1.0 and 0.0, as done in IterativeSolvers.jl
	fill!(p, 0.0)

	i = 0
	while norm_res > tol && i < maxiter
		# precondition - z = "A^{-1}" r
		M!(z, r)

		# update ρ = r^T z
		ρ_prev = ρ
		ρ = dot(z, r)

		# update β and direction p
		β = ρ / ρ_prev
		p .= z .+ β.*p

		# store A*p in z
		z .= p
		mul!(tmp_m, J, p)
		mul!(z, J', tmp_m, 1.0, μ)
		# gemv!('N', 1.0, J, p, 0.0, tmp_m)
		# gemv!('T', 1.0, J, tmp_m, μ, z)

		# update α = r^T z / (p^T A p)
		α = ρ / dot(p, z)

		# update solution
		# x = x + α*p
		# r = r - α*A*p
		axpy!(α, p, x)
		axpy!(-α, z, r)

		norm_res = norm(r)

		i += 1
	end

	flag = 0
	if i == maxiter
		flag = 1
	end

	return flag, i
end

function proj_precondition!(z, r, μ, U, Σ, rank, tmp_m)
	z .= r
	kgemv!('T', rank, 1.0, U, r, 0.0, tmp_m)
	@fastmath @inbounds @simd for j in 1:rank
		tmp_m[j] *= Σ[j]*Σ[j] / (μ + Σ[j]*Σ[j])
	end
	kgemv!('N', rank, -1/μ, U, tmp_m, 1/μ, z)

	return z
end

function no_precondition(z, r)
	z .= r

	return z
end

function retract!(cval, xnew, c!, xtilde, x, method::ProjPenalty{F}) where F

	#= performs primal penalty minimization of
		1/2 || c(z) ||^2 + μ/2 || z - xtilde ||^2
	as μ → 0

	gradient is given by Jct c + μ (z - xtilde)
	uses Gauss-Newton Hessian of the form Jct Jc + μ I = U Σ^2 U' + μ I

	INPUT
	cval - contraint value vector
	xnew - vector in which to store the result
	c! - constraint function of the form c!(cval, x) where y is overwritten
	xtilde - current x
	x - previous iterate
	jac! - function to calculate Jacobian of the form jac!(J, cval, x)
	U, Σ, Vt - SVD decomposition of Jacobian at xtilde
	rank - SVD rank
	μ0 - initial penalty strength μ0
	tol - function tolerance for termination
	maxiter - maximum number of (outer) iterations
	maxiter_pcg - maximum number of (inner) pcg iterations
	work - ProjPenaltyWork struct for storage

	OUTPUT
	flag - 0 = success, 1 = maxiter reached, 2 = pcg maxiter reached, 3 = linesearch failed
	i - total number of (outer) iterations taken
	pcg_iter_count - total cumulative number of (inner) pcg iterations taken

	OVERWRITTEN
	cval, xnew

	=#

	# unpack
	jac! = method.jac!
	U = method.U
	Σ = method.Σ
	Vt = method.Vt
	rank = method.rank
	μ0 = method.μ0
	tol = method.tol
	maxiter = method.maxiter
	maxiter_pcg = method.maxiter_pcg
	work = method.work
	idecomp = method.idecomp
	idata = method.idata

	# unpack work variables
	J = work.J
	tmp_m = work.tmp_m
	r = work.r
	p = work.p
	z = work.z
	dx = work.dx
	g = work.g
	cvalaug = work.cvalaug

	# matrix to use for multiplication
	fulljac = method.ineq ? idecomp' : J

	flag = 0 # return flag

	# initialize x vectors and calculate constraint functions
	xnew .= xtilde
	μ = μ0

	# store length of "x" part
	n = method.ineq ? length(xnew) ÷ 2 : length(xnew)
	m = length(Σ)

	i = 0
	pcg_iter_count = 0
	while i < maxiter
		# calculate new constraint function values and Jacobian
		jac!(J, cval, view(xnew, 1:n))
    	curtol = norm(cval, Inf)

	    if method.ineq
	    	inequality_gradient!(idecomp, xnew, idata)

	    	# update c Jacobian in idecomp
	    	transpose!(idecomp.Jct, J)

	    	# calculate inequality constraint norm
	    	calculate_h!(cvalaug, xnew, idata)

	    	curtol = max(curtol, norm(cvalaug, Inf))
	    end

	    # fill last m values of cvalaug with cval
	    cvalaug[end-m+1:end] .= cval

		# check if tolerance met
		if curtol < tol
			break
		end

	    # calculate right-hand side g = (J' c + μ*(xnew - xtilde))
	    g .= xnew .- xtilde

	    prev_obj_val = dot(cvalaug, cvalaug) + μ*dot(g, g)	# for use below

	    # gemv!('T', 1.0, fulljac, cvalaug, μ, g)
	    mul!(g, fulljac', cvalaug, 1.0, μ)
	    fill!(dx, 0.0)
	    r .= g

	    # do Newton step = (J'J + μI) \ r
	    # M! = (z, r) -> proj_precondition!(z, r, μ, U, Σ, rank, tmp_m)
	    pcg_flag, pcg_i = pcg!(μ, fulljac, no_precondition, dx, r, p, z, tmp_m, tol, maxiter_pcg)
	    pcg_iter_count += pcg_i
	    if pcg_flag > 0
	    	# do no further iterations if pcg exceeds the maximum iteration count
	    	flag = 2
	    	break
	    end

	    # do Armijo backtracking linesearch
	    p .= xnew
	    ar_dot = -dot(g, dx)			# step*gradient

	    α = 1.0

	    xnew .-= α.*dx
	    g .= xnew .- xtilde
	    dist2 = dot(g, g)
	    c!(cval, view(xnew, 1:n))

	    # update inequality constraint function values
	    if method.ineq
	    	calculate_h!(cvalaug, xnew, idata)
	    end

	    cvalaug[end-m+1:end] .= cval

	    armijo_count = 0
	    while dot(cvalaug, cvalaug) + μ*dist2 > prev_obj_val + 1e-4*α*ar_dot
	    	# update values with new step size
	    	α /= 2
	    	xnew .= p .- α.*dx
		    g .= xnew .- xtilde
		    dist2 = dot(g, g)

		    c!(cvalaug, view(xnew, 1:n))

		    # update inequality constraint function values
		    if method.ineq
		    	calculate_h!(cvalaug, xnew, idata)
		    end

	    	cvalaug[end-m+1:end] .= cval

		    # break with error if Armijo iteration count is excessive
	    	armijo_count += 1

	    	if armijo_count == 100
	    		flag = 3
	    		break
	    	end
	    end

	    # update iteration count and μ
		i += 1
		μ = min(μ*0.1, norm(cvalaug))
		# μ *= 0.1
	end

	if i == maxiter
		flag = 1
	end

	return flag, i, pcg_iter_count

end

# ---------------------------- Inequality retractions ---------------------------------------

# retract onto the inequality constraints (y) first and overwrite results in xnewaug
# xnewaug - "xtilde" to be overwritten
# xaug - point at which retraction is based
# idata
function y_retract!(xnewaug, xaug, idata::InequalityData)
	n = length(xaug) ÷ 2

	xnew = view(xnewaug, 1:n)
	ynew = view(xnewaug, n+1:2*n)
	x = view(xaug, 1:n)
	y = view(xaug, n+1:2*n)

	for i in 1:n
		if idata.isline[i]
			# don't need to do anything since step should be in tangent plane
			# but can update to avoid numerical drift
			xnew[i] = ynew[i]
		elseif idata.isparabola[i]
			# retract back to parabola with a second-order retractor
			s = idata.s[i]
			r = idata.r[i]

			g = (-s, -2(y[i] - r))			# gradient at x
			ng = norm(g)
			ux = x[i] - xnew[i] + g[1]/ng
			uy = y[i] - ynew[i] + g[2]/ng

			# solve quadratic and calculate step length, γ, back to parabola
			a = s*uy^2
			b = ux + 2*s*(ynew[i] - r)*uy
			c = xnew[i] + s*(ynew[i] - r)^2 - r

			a1 = -b/(2*a)
			a2 = sqrt(b^2 - 4*a*c)/(2*a)

			γ = min(a1 + a2, a1 - a2)

			# update xnew and ynew
			xnew[i] += γ*ux
			ynew[i] += γ*uy
		else # iscircle
			# retract to circle via simple projection
			c = idata.r[i]
			ρ = sqrt(idata.t[i])

			dist = sqrt((xnew[i] - c)^2 + (ynew[i] - c)^2)

			ynew[i] = c + ρ*(ynew[i] - c)/dist
			xnew[i] = c + ρ*(xnew[i] - c)/dist
		end
	end

	return xnew
end


text:
@testset "Node/HashNode/DynamicNode/DynamicHashNode" begin
    for node in (Node{LeafNode{Float64, 2, 2}, 2, 2}(),
                 HashNode{LeafNode{Float64, 2, 2}, 2, 2}(),
                 DynamicNode{LeafNode{Float64, 2, 2}, 2}(4, 4),
                 DynamicHashNode{LeafNode{Float64, 2, 2}, 2}(4, 4))
        @test TreeArrays.childtype(node) == LeafNode{Float64, 2, 2}
        @test TreeArrays.leaftype(node) == LeafNode{Float64, 2, 2}
        @test TreeArrays.leafeltype(node) == Float64
        @test @inferred(size(node)) == (4, 4)

        @test node[1] === TreeArrays.null(TreeArrays.childtype(node))
        childnull = TreeArrays.null(TreeArrays.childtype(node))

        # cannot set null
        @test_throws Exception node[1] = childnull

        # setindex!
        node[1] = LeafNode{Float64, 2, 2}()
        @test TreeArrays.isactive(node, 1) == true
        node[1][1] = 3
        @test TreeArrays.isactive(node[1], 1) == true

        # setindex! nothing
        node[1] = nothing
        @test isactive(node, 1) == false
        @test isallocated(node, 1) == true

        # allocate!
        TreeArrays.allocate!(node, 1)
        @test isactive(node, 1) == true
        @test isallocated(node, 1) == true
        @test !any(TreeArrays.getmask(node[1]))           # child is still deactivated
        @test TreeArrays.unsafe_getindex(node[1], 1) == 3 # but the value is not changed
        node[1][1] = 3
        @test TreeArrays.isactive(node[1], 1)
        @test node[1][1] == 3
        #
        TreeArrays.allocate!(node, 3)
        @test isactive(node, 3) == true
        @test isallocated(node, 3) == true

        # nleaves
        node[3][1] = 10
        @test TreeArrays.isactive(node[3], 1)
        @test nleaves(node) == 2

        # deactivate!
        TreeArrays.deactivate!(node)
        @test !any(TreeArrays.getmask(node))
        # children are also deactivated
        @test !any(TreeArrays.getmask(TreeArrays.unsafe_getindex(node, 1)))
        @test !any(TreeArrays.getmask(TreeArrays.unsafe_getindex(node, 3)))
        # but still allocated
        @test isallocated(node, 1) == true
        @test isallocated(node, 3) == true
        # other children are not allocated
        @test isallocated(node, 2) == false
        @test isallocated(node, 4) == false

        # cleanup!
        TreeArrays.cleanup!(node)
        for i in eachindex(node)
            @test node[i] === childnull
        end
    end
end


text:
@testset "varinfo.jl" begin
    @testset "TypedVarInfo" begin
        @model gdemo(x, y) = begin
            s ~ InverseGamma(2,3)
            m ~ truncated(Normal(0.0, sqrt(s)), 0.0, 2.0)
            x ~ Normal(m, sqrt(s))
            y ~ Normal(m, sqrt(s))
        end
        model = gdemo(1.0, 2.0)

        vi = VarInfo()
        model(vi, SampleFromUniform())
        tvi = TypedVarInfo(vi)

        meta = vi.metadata
        for f in fieldnames(typeof(tvi.metadata))
            fmeta = getfield(tvi.metadata, f)
            for vn in fmeta.vns
                @test tvi[vn] == vi[vn]
                ind = meta.idcs[vn]
                tind = fmeta.idcs[vn]
                @test meta.dists[ind] == fmeta.dists[tind]
                @test meta.orders[ind] == fmeta.orders[tind]
                @test meta.gids[ind] == fmeta.gids[tind]
                for flag in keys(meta.flags)
                    @test meta.flags[flag][ind] == fmeta.flags[flag][tind]
                end
                range = meta.ranges[ind]
                trange = fmeta.ranges[tind]
                @test all(meta.vals[range] .== fmeta.vals[trange])
            end
        end
    end
    @testset "Base" begin
        # Test Base functions:
        #   string, Symbol, ==, hash, in, keys, haskey, isempty, push!, empty!,
        #   getindex, setindex!, getproperty, setproperty!
        csym = gensym()
        vn1 = @varname x[1][2]
        @test string(vn1) == "x[1][2]"
        @test Symbol(vn1) == Symbol("x[1][2]")

        vn2 = @varname x[1][2]
        @test vn2 == vn1
        @test hash(vn2) == hash(vn1)
        @test inspace(vn1, (:x,))

        function test_base!(vi)
            empty!(vi)
            @test getlogp(vi) == 0
            @test get_num_produce(vi) == 0

            vn = @varname x
            dist = Normal(0, 1)
            r = rand(dist)
            gid = Selector()

            @test isempty(vi)
            @test ~haskey(vi, vn)
            push!(vi, vn, r, dist, gid)
            @test ~isempty(vi)
            @test haskey(vi, vn)

            @test length(vi[vn]) == 1
            @test length(vi[SampleFromPrior()]) == 1

            @test vi[vn] == r
            @test vi[SampleFromPrior()][1] == r
            vi[vn] = [2*r]
            @test vi[vn] == 2*r
            @test vi[SampleFromPrior()][1] == 2*r
            vi[SampleFromPrior()] = [3*r]
            @test vi[vn] == 3*r
            @test vi[SampleFromPrior()][1] == 3*r

            empty!(vi)
            @test isempty(vi)
            push!(vi, vn, r, dist, gid)

            function test_inspace()
                space = (:x, :y, @varname(z[1]), @varname(M[1:10, :]))

                @test inspace(@varname(x), space)
                @test inspace(@varname(y), space)
                @test inspace(@varname(x[1]), space)
                @test inspace(@varname(z[1][1]), space)
                @test inspace(@varname(z[1][:]), space)
                @test inspace(@varname(z[1][2:3:10]), space)
                @test inspace(@varname(M[[2,3], 1]), space)
                @test inspace(@varname(M[:, 1:4]), space)
                @test inspace(@varname(M[1, [2, 4, 6]]), space)
                @test !inspace(@varname(z[2]), space)
                @test !inspace(@varname(z), space)
            end
            test_inspace()
        end
        vi = VarInfo()
        test_base!(vi)
        test_base!(empty!(TypedVarInfo(vi)))
    end
    @testset "flags" begin
        # Test flag setting:
        #    is_flagged, set_flag!, unset_flag!
        function test_varinfo!(vi)
            vn_x = @varname x
            dist = Normal(0, 1)
            r = rand(dist)
            gid = Selector()

            push!(vi, vn_x, r, dist, gid)

            # del is set by default
            @test !is_flagged(vi, vn_x, "del")

            set_flag!(vi, vn_x, "del")
            @test is_flagged(vi, vn_x, "del")

            unset_flag!(vi, vn_x, "del")
            @test !is_flagged(vi, vn_x, "del")
        end
        vi = VarInfo()
        test_varinfo!(vi)
        test_varinfo!(empty!(TypedVarInfo(vi)))
    end
    @testset "setgid!" begin
        vi = VarInfo()
        meta = vi.metadata
        vn = @varname x
        dist = Normal(0, 1)
        r = rand(dist)
        gid1 = Selector()
        gid2 = Selector(2, :HMC)

        push!(vi, vn, r, dist, gid1)
        @test meta.gids[meta.idcs[vn]] == Set([gid1])
        setgid!(vi, gid2, vn)
        @test meta.gids[meta.idcs[vn]] == Set([gid1, gid2])

        vi = empty!(TypedVarInfo(vi))
        meta = vi.metadata
        push!(vi, vn, r, dist, gid1)
        @test meta.x.gids[meta.x.idcs[vn]] == Set([gid1])
        setgid!(vi, gid2, vn)
        @test meta.x.gids[meta.x.idcs[vn]] == Set([gid1, gid2])
    end
    @testset "setval!" begin
        @model function testmodel(x)
            n = length(x)
            s ~ truncated(Normal(), 0, Inf)
            m ~ MvNormal(n, 1.0)
            x ~ MvNormal(m, s)
        end

        x = randn(5)
        model = testmodel(x)

        # UntypedVarInfo
        vi = VarInfo()
        model(vi, SampleFromPrior())

        vicopy = deepcopy(vi)
        DynamicPPL.setval!(vicopy, (m = zeros(5),))
        @test vicopy[@varname(m)] == zeros(5)
        @test vicopy[@varname(s)] == vi[@varname(s)]

        DynamicPPL.setval!(vicopy, (; (Symbol("m[$i]") => i for i in (1, 3, 5, 4, 2))...))
        @test vicopy[@varname(m)] == 1:5
        @test vicopy[@varname(s)] == vi[@varname(s)]

        DynamicPPL.setval!(vicopy, (s = 42,))
        @test vicopy[@varname(m)] == 1:5
        @test vicopy[@varname(s)] == 42

        # TypedVarInfo
        vi = VarInfo(model)

        vicopy = deepcopy(vi)
        DynamicPPL.setval!(vicopy, (m = zeros(5),))
        @test vicopy[@varname(m)] == zeros(5)
        @test vicopy[@varname(s)] == vi[@varname(s)]

        DynamicPPL.setval!(vicopy, (; (Symbol("m[$i]") => i for i in (1, 3, 5, 4, 2))...))
        @test vicopy[@varname(m)] == 1:5
        @test vicopy[@varname(s)] == vi[@varname(s)]

        DynamicPPL.setval!(vicopy, (s = 42,))
        @test vicopy[@varname(m)] == 1:5
        @test vicopy[@varname(s)] == 42
    end
end


text:
using FractionalCalculus
using Test

@testset "Test Direct Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 0, 1.23, 1e-8, RL_Direct())[1], 1.163931646862977; atol = 0.00001)
    @test isapprox(fracint(x->x^4, 0.25, 0, 1.23, 1e-8, RL_Direct())[1], 1.64294134724409; atol = 0.00001)
    @test isapprox(fracint(x->x^4, 0.15, 0, 1.23, 1e-8, RL_Direct())[1], 1.878978485456182; atol = 0.00001)
    @test isapprox(fracint(x->sin(x), 0.5, 0, 0.97, 1e-8, RL_Direct())[1], 0.6442601778843754; atol = 0.00001)
    @test isapprox(fracint(x->sin(x), 0.25, 0, 0.97, 1e-8, RL_Direct())[1], 0.7448298999485293; atol = 0.00001)
    @test isapprox(fracint(x->sin(x), 0.15, 0, 0.97, 1e-8, RL_Direct())[1], 0.7802555196986999; atol = 0.00001)
end

@testset "Test Piecewise Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-4, RL_Piecewise()), 1.163931646862977; atol = 1e-5)
    @test isapprox(fracint(x->x^4, 0.25, 1.23, 1e-4, RL_Piecewise()), 1.64294134724409; atol = 1e-5)
    @test isapprox(fracint(sin, 0.5, 0.97, 1e-4, RL_Piecewise()), 0.6442601778843754; atol = 1e-5)

    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 1e-4, RL_Piecewise())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-3)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-3)
    @test isapprox(result[3], 64.3280200256207; atol = 1e-3)
end

@testset "Test RLInt_Approx" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-4, RLInt_Approx()), 1.163931646862977; atol = 1e-4)
end

@testset "Test RL_LinearInterp" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-4, RL_LinearInterp()), 1.163931646862977; atol = 1e-4)

    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 1e-4, RL_LinearInterp())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-3)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-3)
    @test isapprox(result[3], 64.3280200256207; atol = 1e-3)
end

@testset "Test Simpson Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-7, RLInt_Simpson()), 1.163931646862977; atol = 1e-2)

    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 1e-7, RLInt_Simpson())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-2)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-2)
    @test isapprox(result[3], 64.3280200256207; atol = 1)
end

@testset "Test Trapezoidal Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-7, RLInt_Trapezoidal()), 1.163931646862977; atol = 1e-4)
#=
    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 1e-7, RLInt_Trapezoidal())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-3)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-3)
    @test isapprox(result[3], 64.3280200256207; atol = 1e-3)
    =#
end

@testset "Test Rectangular Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 1e-6, RLInt_Rectangular()), 1.163931646862977; atol = 1e-3)
end

@testset "Test Cubic Spline Interpolation Fractional Integral" begin
    @test isapprox(fracint(x->x^4, 0.5, 1.23, 0.001, RLInt_Cubic_Spline_Interp()), 1.163931646862977; atol = 1e-5)
    @test isapprox(fracint(sin, 0.5, 0.97, 0.01, RLInt_Cubic_Spline_Interp()), 0.6442601778843754; atol = 1e-6)

    result = fracint(x->x^4, 0.5, [1.23, 2, 3], 0.01, RLInt_Cubic_Spline_Interp())
    @test isapprox(result[1], 1.163931646862977; atol = 1e-3)
    @test isapprox(result[2], 10.375035632494745; atol = 1e-3)
    @test isapprox(result[3], 64.3280200256207; atol = 1e-3)
end

@testset "Test fractional integral macros" begin
    @test isapprox(@fracint(x->x^4, 0.5, 1.23), 1.163931646862977; atol = 1e-4)
    @test isapprox(@semifracint(x->x^4, 1.23), 1.163931646862977; atol = 1e-4)
end


text:
#================================================

AbstractSkipList API definition, typedefs, and
shared code for child types.

=================================================#

#===========================
Typedefs
===========================#

"""
    SkipList{T,M} <: AbstractSkipList{T,M}

A non-concurrent skip list. `T` is the type of the element that can be stored in the
`SkipList`, while `M` is the list's mode:

- `M == :List`: the same key can be stored in a skip list multiple times.
- `M == :Set`: only one copy of each key is allowed in the skip list. `SkipListSet{T}` is
  an alias for `SkipList{T,:Set}`.

# Examples

```jldoctest; setup = :(using SkipLists)
julia> list = SkipList{Int64}();

julia> insert!(list, 1); insert!(list, 2); insert!(list, 3);

julia> collect(list)
3-element Array{Int64,1}:
 1
 2
 3
```

"""
mutable struct SkipList{T,M} <: AbstractSkipList{T,M}
    # height_p: the parameter used by the geometric distribution for which the height
    # of a newly-created node is sampled. See random_height for more information.
    height_p::Float64

    # The maximum height of a node in the skip list
    max_height::Int64

    # The maximum number of elements that can be stored in a single SkipList node
    node_capacity::Int64

    # predecessor / successor buffers used to store the results found by the find_node
    # function
    predecessor_buffer::Vector{Node{T,M}}
    successor_buffer::Vector{Node{T,M}}

    # left / right sentinel nodes marking the left and right ends of the skip list
    left_sentinel::Node{T,M}
    right_sentinel::Node{T,M}

    # The height of the skip list. This should be an upper bound on the height of all
    # nodes within the skip list
    height::Int64

    # The number of elements in the skip list
    length::Int64
end

"""
    SkipListSet{T}

A non-concurrent skip list that only allows one copy of each key. `SkipListSet{T}` is an alias for
`SkipList{T,:Set}`.
"""
SkipListSet{T} = SkipList{T,:Set}

#===========================
Shared constructors
===========================#

for list_type in (:SkipList,)
    @eval $list_type{T}(args...; kws...) where T = $list_type{T,:List}(args...; kws...)
end

#===========================
Shared AbstractSkipList API
===========================#

function Base.show(io::IO, list::AbstractSkipList)
    write(io, string(list))
end

function Base.string(list::AbstractSkipList{T,M}) where {T,M}
    type_str = if M == :Set
        "SkipListSet{$T}"
    else
        "SkipList{$T}"
    end

    "$type_str(length = $(length(list)), height = $(height(list)))"
end

"""
    max_height(list::AbstractSkipList)

Return the maximum height that a node can attain in the skip list.
"""
max_height(list::AbstractSkipList) = list.max_height

"""
    height_p(list::AbstractSkipList)

Return the `height_p` parameter used by the skip list to sample a random height
for a newly generated node.
"""
height_p(list::AbstractSkipList) = list.height_p

Base.eltype(::AbstractSkipList{T}) where T = T


text:
####################################################################################################
# Produces required plots for each scenario with appropriate lambda values over a long simulation ##
####################################################################################################

include("../../src/discrete_event_simulation_project.jl")
function results(max_time)
    scenarios = get_scenarios()

    plot_simulation_summary(scenarios[1],
                            max_time = max_time,
                            lambda_range = LinRange(0.1, 5, 20),
                            scenario_label = "scenario 1",
                            save_folder = "test/project_output/scenario_1")

    plot_simulation_summary(scenarios[2],
                            max_time = max_time,
                            lambda_range = LinRange(0.1, 5, 20),
                            scenario_label = "scenario 2",
                            save_folder = "test/project_output/scenario_2")

    plot_simulation_summary(scenarios[3],
                            max_time = max_time,
                            lambda_range = LinRange(0.1, 10, 20),
                            scenario_label = "scenario 3",
                            save_folder = "test/project_output/scenario_3")

    plot_simulation_summary(scenarios[4],
                            max_time = max_time,
                            lambda_range = LinRange(0.1, 0.9, 20),
                            scenario_label = "scenario 4",
                            save_folder = "test/project_output/scenario_4")

    plot_simulation_summary(scenarios[5],
                            max_time = max_time,
                            lambda_range = LinRange(0.1, 1, 20),
                            scenario_label = "scenario 5",
                            save_folder = "test/project_output/scenario_5")

end

results(10^7)


text:
export InsertDeleteMove

struct InsertDeleteMove <: InvMcmcMove
    ν::Int
    len_dist::DiscreteUnivariateDistribution
    counts::Vector{Int}
    function InsertDeleteMove(
        ;ν::Int=1,
        len_dist::DiscreteUnivariateDistribution=Geometric(0.9)
        )
        new(
            ν, len_dist,
            # zeros(Int, ν), zeros(Int, ν), [0,0],
            [0,0]
        )
    end
end

Base.show(io::IO, x::InsertDeleteMove) = print(io, "InsertDeleteMove(ν=$(x.ν), len_dist=$(x.len_dist))")

function multi_insert_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    move::InsertDeleteMove,
    pointers::InteractionSequence{Int},
    ε::Int,
    V::UnitRange
    )
    len_dist = move.len_dist
    log_ratio = 0.0
    # ind = move.ind_add

    n = length(S_prop)+ε
    k = ε
    i = 0
    # j = 0
    while k > 0
        u = rand()
        q = (n - k) / n
        while q > u  # skip
            i += 1
            n -= 1
            q *= (n - k) / n
        end
        i += 1
        # j += 1
        # Insert path at index i
        # i is now index to insert
        # ind[j] = i
        tmp = popfirst!(pointers) # Get storage
        m = rand(len_dist)  # Sample length
        resize!(tmp, m) # Resize
        sample!(V, tmp) # Sample new entries uniformly
        @inbounds insert!(S_prop, i, tmp) # Insert path into S_prop
        log_ratio += - logpdf(len_dist, m) + m*log(length(V))
        n -= 1
        k -= 1
    end
    return log_ratio

end

function multi_delete_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    move::InsertDeleteMove,
    pointers::InteractionSequence{Int},
    ε::Int,
    V::UnitRange
    )

    len_dist = move.len_dist
    log_ratio = 0.0
    # ind = move.ind_del # Store which entries were deleted

    n = length(S_curr)
    k = ε
    i = 0
    # j = 0
    live_index = 0
    while k > 0
        u = rand()
        q = (n - k) / n
        while q > u  # skip
            i += 1
            n -= 1
            q *= (n - k) / n
        end
        i+=1
        # j+=1
        # Delete path
        # i is now index to delete
        # @inbounds ind[j] = i
        @inbounds tmp = popat!(S_prop, i - live_index)
        pushfirst!(pointers, tmp)
        m = length(tmp)
        log_ratio += logpdf(len_dist, m) - m * log(length(V))
        live_index += 1
        n -= 1
        k -= 1
    end
    if length(S_curr)-ε<1
        log_ratio += -Inf
    end
    return log_ratio

end

function prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    move::InsertDeleteMove,
    pointers::InteractionSequence{Int},
    V::UnitRange
    )

    N = length(S_curr)
    ν = move.ν
    ε = rand(1:ν)
    d = rand(0:min(ε, N))
    a = ε - d
    # move.num_add_del[1] = a
    # move.num_add_del[2] = d

    log_ratio = 0.0
    # @show d, a, S_curr, S_prop
    if d > 0
        log_ratio += multi_delete_prop_sample!(
            S_curr, S_prop,
            move, pointers,
            d, V
        )
    end
    if a > 0
        log_ratio += multi_insert_prop_sample!(
            S_curr, S_prop,
            move, pointers,
            a, V
        )
    end

    M = N - d + a
    log_ratio += log(min(ε,N)+1) - log(min(ε, M)+1)

    return log_ratio
end

# function enact_accept!(
#     S_curr::InteractionSequence{Int},
#     S_prop::InteractionSequence{Int},
#     pointers::InteractionSequence{Int},
#     move::InsertDeleteMove
#     )
#     ind = view(move.ind_del, 1:move.num_add_del[2])
#     for i in Iterators.reverse(ind) # If we didnt do reverse would have to update indices
#         @inbounds tmp = popat!(S_curr, i)
#         pushfirst!(pointers, tmp)
#     end
#     ind = view(move.ind_add, 1:move.num_add_del[1])
#     for i in ind
#         tmp = popfirst!(pointers)
#         copy!(tmp, S_prop[i])
#         insert!(S_curr, i, tmp)
#     end
# end

# function enact_reject!(
#     S_curr::InteractionSequence{Int},
#     S_prop::InteractionSequence{Int},
#     pointers::InteractionSequence{Int},
#     move::InsertDeleteMove
#     )
#     ind = view(move.ind_add, 1:move.num_add_del[1])
#     for i in Iterators.reverse(ind)
#         @inbounds tmp = popat!(S_prop, i)
#         pushfirst!(pointers, tmp)
#     end
#     ind = view(move.ind_del, 1:move.num_add_del[2])
#     for i in ind
#         @inbounds tmp = popfirst!(pointers)
#         copy!(tmp, S_curr[i])
#         insert!(S_prop, i, tmp)
#     end
# end


text:
using NanStats
using Printf

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

vec1=[NaN,3,4,NaN]
mn=nanmean(vec1);
varn=nanvar(vec1);
kurn=nankurt(vec1);
sken=nanskew(vec1);

print("vec1=$(vec1)\n");

@testset "nanmean" begin
	@test mn == 3.5
        print("The mean for vec1 ignoring NaN values is $(mn)\n");
end

@testset "nanvar" begin
	@test varn == 0.25
        print("The variance for vec1 ignoring NaN values is $(varn)\n")
end

@testset "nankurt" begin
	@test kurn == 1.0
        print("The kurtosis for vec1 ignoring NaN values is $(kurn)\n")
end

@testset "nanskew" begin
	@test sken == 0.0
        print("The skewness for vec1 ignoring NaN values is $(sken)\n")
end


text:
using BlockArrays

function _concatenate_functions(funcs...)
    function output_function(v)
        out_v = similar(v)
        for (func_ind, func) in enumerate(funcs)
            out_v[Block(func_ind)] = func(v[Block(func_ind)])
        end
        return out_v
    end
end

max_stars_repo_name: f-t-s/CMD

max_issues_repo_name: f-t-s/CMD

max_forks_repo_name: f-t-s/CMD