text:
module Bindings

using Pkg
using LinearAlgebra, LowRankApprox, Statistics, StaticArrays
using JuLIP
using ACE, ACEtb
using ACEtb.Bonds: BondCutoff, get_env, get_env_j, get_env_neighs, get_env_neighs_j, get_all_neighs, get_i_neighs, get_i_neighs_j, eval_bond, get_basis
using ACEtb.SlaterKoster
import ACEtb.SlaterKoster.CodeGeneration
using ACEtb.SlaterKoster: sk2cart
using ACEtb.Utils: read_json, write_json, h5read_SK
using ACEtb.Predictions: predict, train_and_predict
using ACEtb.TBhelpers
using ProgressMeter
using Random
using JSON
using HDF5

bohr2ang = 0.188972598857892E+01
hartree2ev = 27.211386024367243

export set_model, model_predict, acetb_greetings

acetb_dct = Dict()

Bondint_table = nothing
cutoff_func = nothing
saveh5_HH = nothing
saveh5_SS = nothing
saveh5_satoms = nothing
model_onsite_vals = nothing

function buildHS_test(H, S, istart, iend, natoms, coords, species, nnei, inei, ipair, norbs,
                      cutoff, cell, HH, SS, supercell_atoms; MPIproc=1, onsite_HS=nothing)
    if(MPIproc == 1)
       Nprg = iend-istart+1
       if isa(stdout, Base.TTY)
          prgres = Progress(Nprg, dt=0.25, desc="[ Info: |    Calculating ... ",
                            barglyphs=BarGlyphs('|','█', ['▁' ,'▂' ,'▃' ,'▅' ,'▆', '▇'],' ','|',),
                            barlen=20)
       end
    end
    invcell = inv(cell)
    mesh = [9, 9, 9]
    central_atom = prod(mesh) ÷ 2 + 1 # atom 365 for the 9×9×9=729 atom dataset
    rcn = get_tbcells(supercell_atoms, invcell, mesh, central_atom)

    if onsite_HS !== nothing
       atoms = set_JuLIP_atoms(acetb_dct["elm_names"], natoms, coords, species, cell)
       onsite_HS = onsite_HS(atoms)
    end

    for ia = istart:iend
       isp = species[ia]
       offset = ia == 1 ? 0 : sum(nnei[1:ia-1])
       # Onsite blocks
       io = ipair[offset + ia] + 1
       nno = norbs[isp] * norbs[isp]
       shift = convert(Array{Int64,1}, [0,0,0])
       cid = find_row_in_matrix(shift, rcn)
      #  @show HH[cid,:,:]
      #  @show SS[cid,:,:]
       if onsite_HS !== nothing
          os_H, os_S = onsite_HS(ia)
          H[io : io + nno - 1] = os_H
          S[io : io + nno - 1] = os_S
       else
          H[io : io + nno - 1] = vcat(HH[cid,:,:]...) / hartree2ev
          S[io : io + nno - 1] = vcat(SS[cid,:,:]...) / hartree2ev
       end

       # Offsite blocks
       for nj = 1:nnei[ia]
          jn = offset + ia + nj
          ja = inei[jn]
          jsp = species[ja]
          ix = ipair[jn]
          iy = ix + norbs[isp] * norbs[jsp]
          ix += 1
          Rij =  transpose(coords[:,ja] - coords[:,ia])
          if cutoff < norm(Rij)
             continue
          end

          shift = convert(Array{Int64,1}, round.(invcell * Rij'))
          shift = wrap_shift(shift, mesh)
          cid = find_row_in_matrix(shift, rcn)
          H[ix : iy] = vcat(HH[cid,:,:]...)
          S[ix : iy] = vcat(SS[cid,:,:]...)
       end
       if(MPIproc == 1)
          next!(prgres)
       end
    end
    if(MPIproc == 1)
       flush(stdout)
    end
end

function buildHS(SKH_list, H, S, istart, iend, natoms, coords, cell, species, nnei, inei, ipair, i2a, norbs,
                 onsite_HS, Bondint_table, cutoff_func, cutoff; MPIproc=1)

    WriteAllow = false
    if(MPIproc == 1)
       if Threads.threadid() == 1
          WriteAllow = true
       end
       Nprg = iend-istart+1
       if isa(stdout, Base.TTY)
          prgres = Progress(Nprg, dt=0.25, desc="[ Info: |    Calculating ... ",
                            barglyphs=BarGlyphs('|','█', ['▁' ,'▂' ,'▃' ,'▅' ,'▆', '▇'],' ','|',),
                            barlen=20)
       end
    end

    atoms = set_JuLIP_atoms(acetb_dct["elm_names"], natoms, coords, species, cell)
    onsite_HS = onsite_HS(atoms)

    Threads.@threads for ia = istart:iend
       isp = species[ia]
      #  @show ia, isp
       offset = ia == 1 ? 0 : sum(nnei[1:ia-1])

       # Onsite blocks
      #  io = ipair[offset + ia] + 1
      #  @show size(onsite_terms[isp])
      #  for ib = 1:norbs[isp]
      #     H[io] = onsite_terms[isp][ib]
      #     S[io] = 1.0
      #     io += norbs[isp] + 1
      #  end
       io = ipair[offset + ia] + 1
       nno = norbs[isp] * norbs[isp]
       os_H, os_S = onsite_HS(ia)
      #  @show os_H, os_S
       H[io : io + nno - 1] = os_H
       S[io : io + nno - 1] = os_S

       # Offsite blocks
       lnb = length(SKH_list[isp].bonds)
       for nj = 1:nnei[ia]
          jn = offset + ia + nj
          ja = inei[jn]
          jsp = species[ja]
          ix = ipair[jn]
          iy = ix + norbs[isp] * norbs[jsp]
          ix += 1
          R0 =  SVector((coords[:,ja] - coords[:,ia])...)
          if cutoff < norm(R0)
             continue
          end

          # Predictions
          Renv = get_env(acetb_dct["julip_atoms"], R0, ia, cutoff_func)
          VV = Bondint_table(R0,Renv)

          # Set H and S
          E  = sk2cart(SKH_list[isp], R0, VV[1:lnb], WriteAllow=WriteAllow)
          H[ix : iy] = vcat(E...) / hartree2ev # FIXME add units to metadata
          ES = sk2cart(SKH_list[isp], R0, VV[lnb+1:end], WriteAllow=WriteAllow)
          S[ix : iy] = vcat(ES...) / hartree2ev # FIXME add units to metadata
       end
       if(MPIproc == 1)
          if isa(stdout, Base.TTY)
             next!(prgres)
          end
       end
    end
    if(MPIproc == 1)
       flush(stdout)
    end
end

function buildHS_dftb_neigh(SKH_list, H, S, istart, iend, norbs, onsite_terms, Bondint_table, cutoff_func, cutoff; MPIproc=1)

    tbcells = get_translation_cells(acetb_dct["julip_atoms"].cell, cutoff)
    nnei, inei, i2a, icellvec, coords, species = get_neighbours(acetb_dct["julip_atoms"], tbcells, cutoff)
    ipair = get_sparse_indexing(acetb_dct["julip_atoms"], nnei, inei, i2a, norbs)

    inei = vcat(inei...)
    ipair = vcat(ipair...)
    coords = transpose(coords)
    println("---")
    println(size(nnei))
    println(size(inei))
    println(size(ipair))
    println(size(coords))

    WriteAllow = false
    if(MPIproc == 1)
       if Threads.threadid() == 1
          WriteAllow = true
       end
       Nprg = iend-istart+1
       if isa(stdout, Base.TTY)
          prgres = Progress(Nprg, dt=0.25, desc="[ Info: |    Calculating ... ",
                            barglyphs=BarGlyphs('|','█', ['▁' ,'▂' ,'▃' ,'▅' ,'▆', '▇'],' ','|',),
                            barlen=20)
       end
    end

    Rt = get_i_neighs(1, length(acetb_dct["julip_atoms"]), coords, nnei, inei)
    #Rt = get_all_neighs(acetb_dct["natoms"], coords, nnei, inei)
    #Rt = get_i_neighs(1, natoms, coords, nnei, inei)
    #Rt, jt = get_i_neighs_j(istart, iend, coords, nnei, inei)

    Threads.@threads for ia = istart:iend
       isp = species[ia]
       offset = ia == 1 ? 0 : sum(nnei[1:ia-1])
       # Onsite blocks
       io = ipair[offset + ia] + 1
       for ib = 1:norbs[isp]
          H[io] = onsite_terms[isp][ib]
          S[io] = 1.0
          io += norbs[isp] + 1
       end
       lnb = length(SKH_list[isp].bonds)

       # Offsite blocks
       for nj = 1:nnei[ia]
          jn = offset + ia + nj
          ja = inei[jn]
          jsp = species[ja]
          ix = ipair[jn]
          iy = ix + norbs[isp] * norbs[jsp]
          ix += 1
          R0 =  Rt[ia][nj]
          if cutoff < norm(R0)
             continue
          end

          # Predictions
          #Renv = get_env_neighs(vcat(Rt[ia],.-Rt[i2a[ja]]), R0, cutoff_func)
          #Renv = get_env_neighs(Rt[ia], R0, cutoff_func)
          #if(MPIproc == 1)
          #   Renv2, jl2 = get_env_neighs_j(Rt[ia], R0, cutoff_func)
          #   jlist2 = [ i2a[jt[ia][jj]] for jj in jl2 ]
          #end
          #Renv, jlist = get_env_j(acetb_dct["julip_atoms"], R0, ia, cutoff_func)
          #if(MPIproc == 1)
          #   for j1 in jlist
          #      if j1 ∉ jlist2
          #         println("ia: ",ia," ja: ",ja," j1: ",j1)
          #      end
          #   end
          #   for j2 in jlist2
          #      if j2 ∉ jlist
          #          println("ia: ",ia," ja: ",ja," j2: ",j2)
          #      end
          #   end
          #end

          Renv = get_env(acetb_dct["julip_atoms"], R0, ia, cutoff_func)
          VV = Bondint_table(R0,Renv)

          # Set H and S
          E  = sk2cart(SKH_list[isp], R0, VV[1:lnb], WriteAllow=WriteAllow)
          H[ix : iy] = vcat(E...)
          ES = sk2cart(SKH_list[isp], R0, VV[lnb+1:end], WriteAllow=WriteAllow)
          S[ix : iy] = vcat(ES...)
       end
       if(MPIproc == 1)
          if isa(stdout, Base.TTY)
             next!(prgres)
          end
       end
    end
    if(MPIproc == 1)
       flush(stdout)
    end
end

function get_list_str(str::Array{UInt8}, ln, stride)
   slen = Int64(ln)
   snames = []
   nstr = String( copy(str[1:slen]) )
   if ln == stride - 1
      push!(snames,strip(nstr[1:ln-1]))
   else
      si = 1
      sf = 0
      for i = 1:ln
         sf += 1
         if sf == stride
            elm = nstr[si:i-1]
            elm = strip(elm)
            push!(snames,elm)
            si = i + 1
            sf = 0
         end
      end
   end
   return snames
end

function get_specie_name(str::Array{UInt8}, ln, n)
   slen = Int64(ln / n) - 1
   snames = []
   for i = 1:n
      si = (i - 1) * slen + (i - 1) + i
      sf = i * slen - 1 + (i - 1)
      elm = String( copy(str[si:sf]) )
      elm = strip(elm)
      push!(snames,elm)
   end
   return snames
end

function set_JuLIP_atoms(elm_names, natoms, pos, species, cell)
   atnums = []
   #atmass = []
   for i=1:natoms
      sym = Symbol(elm_names[species[i]])
      #am = atomic_mass(sym)
      az = atomic_number(sym)
      push!(atnums,az)
      #push!(atmass,am)
   end

   return Atoms(; X = pos[:,1:natoms], Z = atnums, cell = cell,
                pbc = [true, true, true])
end

function set_model(natoms, nspecies,
                   coords::Array{Float64},
                   latvecs::Array{Float64},
                   tcellv::Array{Float64},
                   rcellv::Array{Float64},
                   origin::Array{Float64},
                   species::Array{Int32},
                   nshells::Array{Int32},
                   norbe::Array{Int32},
                   norba::Array{Int32},
                   maxns, maxno, maxs,
                   maxo, toto,
                   angshell::Array{Int32},
                   ishell::Array{Int32},
                   posshell::Array{Int32},
                   lstr, specienames::Array{UInt8},
                   cutoff,
                   lfn, fnames::Array{UInt8},
                   MPIproc,
                   stat_jl::Array{Int32})
    if(MPIproc == 1)
        @info "┌── ACEtb : Julia set_model function."
    end
    try
        global acetb_dct["natoms"] = natoms
        global acetb_dct["ntypes"] = nspecies
        pos = reshape(coords,3,:) ./ bohr2ang
        cell = reshape(latvecs,3,:) ./ bohr2ang
        tcells = reshape(tcellv,3,:) ./ bohr2ang
        rcells = reshape(rcellv,3,:) ./ bohr2ang
        global acetb_dct["latvecs"] = cell[:,:]
        global acetb_dct["tcellv"] = tcells[:,:]
        global acetb_dct["rcellv"] = rcells[:,:]
        global acetb_dct["origin"] = origin[:]
        global acetb_dct["norbe"] = norbe

        elm_names = get_specie_name(specienames, lstr, nspecies)
        global acetb_dct["elm_names"] = elm_names[:]

        at = set_JuLIP_atoms(elm_names, natoms, pos, species, cell)
        global acetb_dct["julip_atoms"] = at

        SKH_list = set_SK_orbitals(nspecies,nshells,angshell)
        global acetb_dct["SKH_list"] = SKH_list

        model_files = get_list_str(fnames, lfn, 200)
        global acetb_dct["model_files"] = model_files

        if(MPIproc == 1)
            @info "│    Reading acetb.json file..."
        end
        filedata = read_json(model_files[1])
        if(MPIproc == 1)
            @info "│    Reading is done."
        end
        global acetb_dct["file_data"] = filedata
        cutoff_params = filedata["model"]["cutoff_params"]
        fit_params = filedata["model"]["fit_params"]
        Bpredict = true
        Bfit = false
        if haskey(filedata["model"], "predict")
            predict_setting = filedata["model"]["predict"]
            if predict_setting == 0
                Bpredict = false
            else
                Bpredict = true
            end
            if haskey(filedata, "training_datasets")
                trdata = filedata["training_datasets"]
            end
        end
        if haskey(filedata["model"], "fit")
            fit_setting = filedata["model"]["fit"]
            if fit_setting == 0
                Bfit = false
            else
                Bfit = true
            end
        end
        if Bpredict
            if Bfit
                if(MPIproc == 1)
                    @info "│    Fitting..."
                end
                Bint_table, cutf_func, train_dict = train_and_predict(trdata, elm_names, cutoff_params, fit_params; MPIproc=MPIproc)
                if(MPIproc == 1)
                    @info "│    Saving potential..."
                    write_json(model_files[1], train_dict)
                end
                global Bondint_table = Bint_table
                global cutoff_func = cutf_func
                if(MPIproc == 1)
                    @info "│    Fitting is done."
                end
            else
                if(MPIproc == 1)
                    @info "│    Loading model..."
                end
                Bint_table, cutf_func = predict(filedata, cutoff_params, fit_params; MPIproc=MPIproc)
                global Bondint_table = Bint_table
                global cutoff_func = cutf_func
                if(MPIproc == 1)
                    @info "│    Model is loaded."
                end
            end
        else
            if(MPIproc == 1)
                @info "│    This mode is test only. It will not use model predictions."
                @info "│    Setting bond integral table..."
            end
            HSfile = filedata["HS_datasets"][1]
            if(MPIproc == 1)
                @info "│    Reading saved H,S file:" HSfile
            end
            HSdata = h5read_SK(HSfile; get_HS=true, get_atoms=true, get_metadata=true, get_energies=true)
            global saveh5_HH = permutedims(HSdata[2][1], [3, 2, 1])
            global saveh5_SS = permutedims(HSdata[2][2], [3, 2, 1])
            global saveh5_satoms = HSdata[6]
            write_extxyz("saveh5_satoms.xyz", saveh5_satoms)

            if(MPIproc == 1)
               @info "│    Setting H,S is done."
            end
        end
        global onsite_vals = filedata["onsite-terms"]
        stat_jl[1] = 0
    catch
        stat_jl[1] = 1
        for (exc, bt) in Base.catch_stack()
            if(MPIproc == 1)
                showerror(stdout, exc, bt)
                println()
            end
        end
    end
    if(MPIproc == 1)
        @info "└── ACEtb : Done in Julia module."
    end
    flush(stdout)
end

function ACE2tb_wrapper_predict_HS(model_ID, atoms; λ=1e-3)
   @show model_ID

   dict_atoms = JuLIP.write_dict(atoms)
   open("structure.json", "w") do f
      write(f, JSON.json(dict_atoms))
   end

   cmd = setenv(`$(joinpath(Sys.BINDIR, Base.julia_exename())) --project=$(homedir())/.julia/dev/ACE2tb compute.jl $(model_ID) $(pwd())/structure.json $(pwd())/onsite_HS.h5`,
                dir="$(homedir())/.julia/dev/ACE2tb/test")
   @show cmd
   run(cmd)

   H, S = h5open("onsite_HS.h5") do f
      read(f, "H"), read(f, "S")
   end

   println("Regularisation λ = $λ")

   return ia -> (reshape(H[:, :, ia], :),
                 reshape(S[:, :, ia] + λ * I, :))
end

function model_predict(iatf, iatl, natoms,
                       coords::Array{Float64},
                       latvecs::Array{Float64},
                       nspecies, nspecarr, species::Array{Int32},
                       nH, H::Array{Float64},
                       nS, S::Array{Float64},
                       nneigh::Array{Int32},
                       ineigh::Array{Int32},
                       ipair::Array{Int32},
                       i2a::Array{Int32}, cutoff,
                       MPIproc,
                       stat_jl::Array{Int32})
    if(MPIproc == 1)
        @info "┌── ACEtb : Julia model_predict function."
    end
    try
        pos = reshape(coords,3,:) ./ bohr2ang
        cell = reshape(latvecs,3,:) ./ bohr2ang
        cutoff /= bohr2ang

        elm_names = acetb_dct["elm_names"]
        #at = set_JuLIP_atoms(elm_names, natoms, pos, species, cell)

        SKH_list = acetb_dct["SKH_list"]
        model_files = acetb_dct["model_files"]
        filedata = acetb_dct["file_data"]

        predict_params = filedata["model"]["predict"]
        onsite_vals = filedata["onsite-terms"]
        norbe = acetb_dct["norbe"]

        if onsite_vals isa AbstractDict
            for isp = 1:nspecies
                sp = elm_names[species[isp]]
                if onsite_vals[sp] isa AbstractVector
                   MPIproc == 1 && @info "│    Using $(norbe[isp]) element diagonal onsite approximation for species $sp"
                   @assert length(onsite_vals[sp]) == norbe[isp]
                   onsite_HS = coords -> (ia -> (reshape(diagm(Float64.(onsite_vals[sp])), :), reshape(1.0*I(norbe[isp]) |> Matrix, :)))
                elseif onsite_vals[sp] isa AbstractDict
                   MPIproc == 1 && @info "│    Using $(norbe[isp]) × $(norbe[isp]) onsite block for species $sp"
                   onsite_HS = coords -> (ia -> (vcat(onsite_vals[sp]["H"]...), vcat(onsite_vals[sp]["S"]...)))
                else
                   error("Invalid onsite specification for species $sp: $(onsite_vals[sp])")
                end
            end
        elseif onsite_vals isa AbstractString
           onsite_HS = coords -> ACE2tb_wrapper_predict_HS(onsite_vals, coords)
        else
           error("Invalid onsite specification $(onsite_vals)")
        end

        MPIproc == 1 && @info "│    Calculating bond integrals..."
        if predict_params == 0
           buildHS_test(H, S, iatf, iatl, natoms, pos, species, nneigh, ineigh, ipair, norbe,
                        cutoff, cell, saveh5_HH, saveh5_SS, saveh5_satoms; MPIproc=MPIproc, onsite_HS=onsite_HS)
        elseif predict_params == -1
           buildHS_dftb_neigh(SKH_list, H, S, iatf, iatl, norbe, onsite_terms, Bondint_table, cutoff_func, cutoff; MPIproc=1)
        else
           buildHS(SKH_list, H, S, iatf, iatl, natoms, pos, cell, species, nneigh, ineigh, ipair, i2a, norbe, onsite_HS,
                   Bondint_table, cutoff_func, cutoff; MPIproc=MPIproc)
        end
        stat_jl[1] = 0
    catch
        stat_jl[1] = 1
        for (exc, bt) in Base.catch_stack()
            if(MPIproc == 1)
                showerror(stdout, exc, bt)
                println()
            end
        end
    end
    if(MPIproc == 1)
        @info "└── ACEtb : Done in Julia module."
    end
    flush(stdout)
end

function acetb_greetings()
   ctx = Pkg.Operations.Context()
   acetb_version = string(ctx.env.manifest[ctx.env.project.deps["ACEtb"]].version)
   acetb_hash = split(string(ctx.env.project.deps["ACEtb"]),"-")[1]
   if isa(stdout, Base.TTY)
      println("  \033[31m\033[1m  ___\033[38;5;208m ___\033[32m ___\033[34m _____\033[38;5;54m\033[1m _    \033[0m       \n"*
              "  \033[31m\033[1m |_  \033[38;5;208m|  _\033[32m| __\033[34m|_   _\033[38;5;54m\033[1m| |_   \033[0m\033[38;5;141m\033[1mVersion\033[0m   \n"*
              "  \033[31m\033[1m | . \033[38;5;208m| |_\033[32m| __|\033[34m | | \033[38;5;54m\033[1m| . |  \033[0m\033[38;5;200m\033[1mv",acetb_version,"\033[0m   \n"*
              "  \033[31m\033[1m |___\033[38;5;208m|___\033[32m|___|\033[34m |__|\033[38;5;54m\033[1m|___|  \033[0m\033[1m[",acetb_hash,"]\033[0m\n")
   else
      println("    ___ ___ ___ _____ _            \n"*
              "   |_  |  _| __|_   _| |_   Version\n"*
              "   | . | |_| __| | | | . |  v",acetb_version,"   \n"*
              "   |___|___|___| |__||___|  [",acetb_hash,"]\n")
   end
   println("   ┌───────────────────────────┐\n"*
           "   │   Developers:             │\n"*
           "   ├───────────────────────────┤\n"*
           "   │     Geneviève Dusson      │\n"*
           "   │     Berk Onat             │\n"*
           "   │     Reinhard Maurer       │\n"*
           "   │     Christoph Ortner      │\n"*
           "   │     James R. Kermode      │\n"*
           "   └───────────────────────────┘\n")
   println("Julia threads:               ",Threads.nthreads())
   flush(stdout)
end

end # End of module


text:
using Distributions

function imcmc_multi_insert_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mcmc::SimMcmcInsertDeleteLengthCentered,
    ind::AbstractVector{Int},
    V::UnitRange,
    K_in_ub::Int
    )

    prop_pointers = mcmc.prop_pointers
    ν_td = mcmc.ν_td
    N = length(S_curr)
    len_dist = truncated(Poisson(mean(length, S_curr)), 1, Inf)

    log_ratio = 0.0
    for i in ind
        tmp = popfirst!(prop_pointers)
        m = rand(len_dist)
        resize!(tmp, m)
        sample!(V, tmp)
        @inbounds insert!(S_prop, i, tmp)
        log_ratio += - logpdf(len_dist, m) + m*log(length(V)) - Inf * (m > K_in_ub)
    end
    log_ratio += log(ν_td) - log(min(ν_td,N))
    return log_ratio

end

function imcmc_multi_insert_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mcmc::SimMcmcInsertDeleteLengthCentered,
    ε::Int,
    V::UnitRange,
    K_in_ub::Int
    )

    prop_pointers = mcmc.prop_pointers
    ν_td = mcmc.ν_td
    len_dist = truncated(Poisson(mean(length, S_curr)), 1, Inf)
    N = length(S_curr)
    log_ratio = 0.0
    ind = mcmc.ind_td

    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!(prop_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)) - Inf * (m > K_in_ub)  # Add to log_ratio term
        n -= 1
        k -= 1
    end
    log_ratio += log(ν_td) - log(min(ν_td,N))
    return log_ratio

end

function imcmc_multi_delete_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mcmc::SimMcmcInsertDeleteLengthCentered,
    V::UnitRange
    )

    prop_pointers = mcmc.prop_pointers
    ν_td = mcmc.ν_td
    N = length(S_curr)
    len_dist = truncated(Poisson(mean(length, S_curr)), 1, Inf)
    log_ratio = 0.0
    ind = mcmc.ind_td

    n = length(S_prop)
    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!(prop_pointers, tmp)
        m = length(tmp)
        log_ratio += logpdf(len_dist, m) - m * log(length(V))
        live_index += 1
        n -= 1
        k -= 1
    end

    log_ratio += log(min(ν_td,N)) - log(ν_td)
    return log_ratio

end

function imcmc_multi_delete_prop_sample!(
    S_curr::InteractionSequence{Int},
    S_prop::InteractionSequence{Int},
    mcmc::SimMcmcInsertDeleteLengthCentered,
    ind::AbstractVector{Int},
    V::UnitRange
    )

    prop_pointers = mcmc.prop_pointers
    ν_td = mcmc.ν_td
    N = length(S_curr)
    len_dist = truncated(Poisson(mean(length, S_curr)), 1, Inf)

    log_ratio = 0.0

    for i in Iterators.reverse(ind)
        @inbounds tmp = popat!(S_prop, i)
        pushfirst!(prop_pointers, tmp)
        m = length(tmp)
        log_ratio += logpdf(len_dist, m) - m * log(length(V))
    end

    log_ratio += log(min(ν_td,N)) - log(ν_td)
    return log_ratio

end

ext: jl


text:
export BTrDBException

struct BTrDBException <: Exception
    msg::String
    BTrDBException() = new("")
    BTrDBException(msg) = new(msg)
end


text:
using Griddly
using Makie
using Documenter: deploydocs, makedocs

makedocs(sitename = "Griddly.jl", modules = [Griddly], doctest = false,
pages = [
        "Home" => "index.md",
        "Library Outline" => Any[
            "Observer" => "lib/observer.md",
            "Grid" => "lib/grid.md",
            "Game Process" => "lib/game_process.md",
            "Player" => "lib/player.md",
            "Render Tools" => "lib/render_tool.md",
        ],
        "Examples" => Any[
            "MiniGridDoggo" => "examples/minigrid_doggo.md",
            "Multi Player" => "examples/multiplayer.md"
        ],
    ])
deploydocs(repo = "github.com/luhervier/Griddly.jl.git")


text:
using SeisIO, SeisNoise, DTWDT, PlotlyJS, JLD2
import LinearAlgebra.pinv

# This is a script to test the functionality of the DTW functions of DTWDT.jl
# This assumes synthetic data generated via genSynth.jl

function testDTW(finame::String, foname::String, InputDict::Dict, type::String; plot::Bool=false)
    valData = jldopen(finame)
    dvVlist = valData["info/dvVlist"]
    noiselist = valData["info/noiselist"]

    f_err = jldopen(foname, "a+")
    if type=="realData"
        f_err["info/dvV"] = dvVlist
        f_err["info/noise"] = noiselist
    end

    maxLag = InputDict["maxLag"]
    b = InputDict["b"]
    st = InputDict["st"]
    dt = InputDict["dt"]
    direction = InputDict["direction"]
    mintime = InputDict["mintime"]

    for dvV in dvVlist
        println("dv/v: $dvV")
        for noiselvl in noiselist
            signals = valData["$type/$dvV.$noiselvl"]
            (u1, u2) = signals

            npts = length(u1)
            tvec = (collect(0:npts-1) .* dt) .+ mintime

            (stbarTime, stbar, dist, error) = dtwdt(u1, u2, dt, maxLag=maxLag, b=b, direction=direction)

            # linear regression for dv/V
            A = ones(npts,2)
            A[:,2] = tvec
            coef = pinv(A) * stbarTime
            (a, m) = coef

            if plot==true
                plot_dtt(tvec, stbarTime, m, dvV)
            end

            f_err["$type/$dvV.$noiselvl"] = [m, error]
        end
    end
    close(f_err)
    close(valData)
end

function plot_dtt(tvec::Array{Float64,1}, stbarTime::Array{Float64,1}, m::Float64, dvV::Float64)
    trace1 = PlotlyJS.scatter(;x=tvec, y=stbarTime, name="Predicted dt")
    trace2 = PlotlyJS.scatter(;x=tvec, y=tvec .* m, name="Regression")
    trace3 = PlotlyJS.scatter(;x=tvec, y=tvec .* dvV, name="True dt")
    p=PlotlyJS.plot([trace1, trace2, trace3])
    display(p)
    readline()
end

# MWCS input parameters
InputDict_real = Dict( "maxLag"    => 200,
                       "b"         => 1,
                       "st"        => 0,
                       "dt"        => 0.05,
                       "direction" => 0,
                       "mintime"   => -100.0)

InputDict_synth = Dict( "maxLag"    => 200,
                        "b"         => 1,
                        "st"        => 0,
                        "dt"        => 0.05,
                        "direction" => 0,
                        "mintime"   => 0.0)

finame = "verificationData.jld2"
foname = "verificationDataError_DTW.jld2"
type = "rickerConv"

for type in ["rickerConv", "dampedSinusoid", "realData"]
    if type == "realData"
        InputDict = InputDict_real
    else
        InputDict = InputDict_synth
    end
    testDTW(finame, foname, InputDict, type)
end

max_stars_repo_name: kura-okubo/SeisXcorrelation

max_issues_repo_name: kura-okubo/SeisXcorrelation

max_forks_repo_name: kura-okubo/SeisXcorrelation


text:
#=

Tests for Discrete Decision  Processes (DDP)

Original Python Author: Daisuke Oyama
Authors: Spencer Lyon and Matthew McKay

Tests for markov/ddp.jl

=#

@testset "Testing markov/dpp.jl" begin

    #-Setup-#

    # Example from Puterman 2005, Section 3.1
    beta = 0.95

    # Formulation with Dense Matrices R: n x m, Q: n x m x n
    n, m = 2, 2  # number of states, number of actions
    R = [5.0 10.0; -1.0 -Inf]
    Q = Array{Float64}(n, m, n)
    Q[:, :, 1] = [0.5 0.0; 0.0 0.0]
    Q[:, :, 2] = [0.5 1.0; 1.0 1.0]

    ddp0 = DiscreteDP(R, Q, beta)

    # Formulation with state-action pairs
    L = 3  # Number of state-action pairs
    s_indices = [1, 1, 2]
    a_indices = [1, 2, 1]
    R_sa = [R[1, 1], R[1, 2], R[2, 1]]
    Q_sa = spzeros(L, n)
    Q_sa[1, :] = Q[1, 1, :]
    Q_sa[2, :] = Q[1, 2, :]
    Q_sa[3, :] = Q[2, 1, :]
    ddp0_sa = DiscreteDP(R_sa, Q_sa, beta, s_indices, a_indices)

    # List of ddp formulations
    ddp0_collection = (ddp0, ddp0_sa)

    # Maximum Iteration and Epsilon for Tests
    max_iter = 200
    epsilon = 1e-2

    # Analytical solution for beta > 10/11, Example 6.2.1
    v_star = [(5-5.5*beta)/((1-0.5*beta)*(1-beta)), -1/(1-beta)]
    sigma_star = [1, 1]

    @testset "bellman_operator methods" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp in ddp0_collection
        	@test isapprox(bellman_operator(ddp, v_star), v_star)
    	end
    end

    @testset "RQ_sigma" begin
        nr, nc = size(R)
        # test for DDP
        sigmas = ([1, 1], [1, 2], [2, 1], [2, 2])
        for sig in sigmas
            r, q = RQ_sigma(ddp0, sig)

            for i_r in 1:nr
                @test r[i_r] == ddp0.R[i_r, sig[i_r]]
                for i_c in 1:length(sig)
                    @test vec(q[i_c, :]) == vec(ddp0.Q[i_c, sig[i_c], :])
                end
            end
        end

        # TODO: add test for DDPsa
    end

    @testset "compute_greedy methods" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp in ddp0_collection
        	@test compute_greedy(ddp, v_star) == sigma_star
    	end
    end

    @testset "evaluate_policy methods" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp in ddp0_collection
        	@test isapprox(evaluate_policy(ddp, sigma_star), v_star)
        end
    end

    @testset "methods for subtypes != (Float64, Int)" begin
        float_types = [Float16, Float32, Float64, BigFloat]
        int_types = [Int8, Int16, Int32, Int64, Int128,
                     UInt8, UInt16, UInt32, UInt64, UInt128]

        for ddp in ddp0_collection
            for f in (bellman_operator, compute_greedy)
                for T in float_types
                    f_f64 = f(ddp, [1.0, 1.0])
                    f_T = f(ddp, ones(T, 2))
                    @test isapprox(f_f64, convert(Vector{eltype(f_f64)}, f_T))
                end

                # only Integer subtypes can be Rational type params
                # NOTE: Only the integer types below don't overflow for this example
                for T in [Int64, Int128]
                    @test f(ddp, [1//1, 1//1]) == f(ddp, ones(Rational{T}, 2))
                end
            end

            for T in float_types, S in int_types
                v = ones(T, 2)
                s = ones(S, 2)
                # just test that we can call the method and the result is
                # deterministic
                @test bellman_operator!(ddp, v, s) == bellman_operator!(ddp, v, s)
            end

            for T in int_types
                s = T[1, 1]
                @test isapprox(evaluate_policy(ddp, s), v_star)
            end
        end
    end

    @testset "compute_greedy! changes ddpr.v" begin
        res = solve(ddp0, VFI)
        res.Tv[:] = 500.0
        compute_greedy!(ddp0, res)
        @test maximum(abs, res.Tv - 500.0) > 0
    end

    @testset "value_iteration" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp_item in ddp0_collection
            # Compute Result
            res = solve(ddp_item, VFI)
            v_init = [0.0, 0.0]
            res_init = solve(ddp_item, v_init, VFI; epsilon=epsilon)

            # Check v is an epsilon/2-approxmation of v_star
            @test maximum(abs, res.v - v_star) < epsilon/2
            @test maximum(abs, res_init.v - v_star)    < epsilon/2

            # Check sigma == sigma_star.
            # NOTE we need to convert from linear to row-by-row index
            @test res.sigma == sigma_star
            @test res_init.sigma == sigma_star
        end
    end

    @testset "policy_iteration" begin
        # Check both Dense and State-Action Pair Formulation
        for ddp_item in ddp0_collection
            res = solve(ddp_item, PFI)
            v_init = [0.0, 1.0]
            res_init = solve(ddp_item, v_init, PFI)

            # Check v == v_star
            @test isapprox(res.v, v_star)
            @test isapprox(res_init.v, v_star)

            # Check sigma == sigma_star
            @test res.sigma == sigma_star
            @test res_init.sigma == sigma_star
        end
    end

    @testset "DiscreteDP{Rational,_,_,Rational} maintains Rational" begin
        ddp_rational = DiscreteDP(map(Rational{BigInt}, R),
                                  map(Rational{BigInt}, Q),
                                  map(Rational{BigInt}, beta))
        # do minimal number of iterations to avoid overflow
        vi = Rational{BigInt}[1//2, 1//2]
        @test eltype(solve(ddp_rational, VFI; max_iter=1, epsilon=Inf).v) == Rational{BigInt}
        @test eltype(solve(ddp_rational, vi, PFI; max_iter=1).v) == Rational{BigInt}
        @test eltype(solve(ddp_rational, vi, MPFI; max_iter=1, k=1, epsilon=Inf).v) == Rational{BigInt}
    end

    @testset "DiscreteDP{Rational{BigInt},_,_,Rational{BigInt}}  works" begin
        ddp_rational = DiscreteDP(map(Rational{BigInt}, R),
                                  map(Rational{BigInt}, Q),
                                  map(Rational{BigInt}, beta))
        # do minimal number of iterations to avoid overflow
        r1 = solve(ddp_rational, PFI)
        r2 = solve(ddp_rational, MPFI)
        r3 = solve(ddp_rational, VFI)
        @test maximum(abs, r1.v-v_star) < 1e-13
        @test r1.sigma == r2.sigma
        @test r1.sigma == r3.sigma
        @test r1.mc.p == r2.mc.p
        @test r1.mc.p == r3.mc.p
    end

    @testset "modified_policy_iteration" begin
        for ddp_item in ddp0_collection
            res = solve(ddp_item, MPFI)
            v_init = [0.0, 1.0]
            res_init = solve(ddp_item, v_init, MPFI)

                    # Check v is an epsilon/2-approxmation of v_star
            @test maximum(abs, res.v - v_star) < epsilon/2
            @test maximum(abs, res_init.v - v_star) < epsilon/2

            # Check sigma == sigma_star
            @test res.sigma == sigma_star
            @test res_init.sigma == sigma_star

            #Test Modified Policy Iteration k0
            k = 0
            res = solve(ddp_item, MPFI; max_iter=max_iter, epsilon=epsilon, k=k)

            # Check v is an epsilon/2-approxmation of v_star
            @test maximum(abs, res.v - v_star) < epsilon/2

            # Check sigma == sigma_star
            @test res.sigma == sigma_star
        end
    end

    @testset "DDPsa constructor" begin
        @testset "feasbile action pair" begin
            _R = [1.0, 0.0, 0.0, 1.0]
            _Q = fill(1/3, 4, 3)
            _s_ind = [1, 1, 3, 3]
            _a_ind = [1, 2, 1, 2]
            @test_throws ArgumentError DiscreteDP(_R, _Q, beta, _s_ind, _a_ind)
        end

        _R, _Q = R_sa, Q_sa
        _s_ind = [1, 1, 2]
        _a_ind = [1, 2, 1]

        @testset "beta in [0, 1)" begin
            @test_throws ArgumentError DiscreteDP(_R, _Q, -eps(), _s_ind, _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, _Q, 1.0, _s_ind, _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, _Q, 1+eps(), _s_ind, _a_ind)
        end

        @testset "argument sizes" begin
            # NQ != 2
            @test_throws ArgumentError DiscreteDP(_R, rand(4, 3, 1), beta, _s_ind, _a_ind)

            # NR != 1
            @test_throws ArgumentError DiscreteDP(rand(4, 1), _Q, beta, _s_ind, _a_ind)

            # incorrect lengths
            @test_throws ArgumentError DiscreteDP(rand(2), _Q, beta, _s_ind, _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, rand(5, 2), beta, _s_ind, _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, _Q, beta, rand(1:3, 2), _a_ind)
            @test_throws ArgumentError DiscreteDP(_R, _Q, beta, _s_ind, rand(1:3, 2))
        end

        @testset "duplicate sa pair" begin
            @test_throws ArgumentError DiscreteDP(_R, _Q, beta, _s_ind, [1, 1, 2])
        end
    end

    @testset "DDP constructor" begin
        @testset "beta in [0, 1)" begin
            @test_throws ArgumentError DiscreteDP(R, Q, -eps())
            @test_throws ArgumentError DiscreteDP(R, Q, 1.0)
            @test_throws ArgumentError DiscreteDP(R, Q, 1+eps())
        end

        @testset "feasbile action pair" begin
            #Dense Matrix
            n, m = 2, 2
            _R = [-Inf -Inf; 1.0 2.0]

            _Q = Array{Float64}(n, m, n)
            _Q[:, :, 1] = [0.5 0.0; 0.0 0.0]
            _Q[:, :, 2] = [0.5 1.0; 1.0 1.0]
            _beta = 0.95

            @test_throws ArgumentError DiscreteDP(_R, _Q, _beta)
        end

        @testset "R, Q sizes" begin
            # NQ != 3
            @test_throws ArgumentError DiscreteDP(R, zeros(2, 2), beta)

            # NR != 2
            @test_throws ArgumentError DiscreteDP(zeros(1), Q, beta)

            # incompatible dimensions
            @test_throws ArgumentError DiscreteDP(zeros(2, 3), Q, beta)
            @test_throws ArgumentError DiscreteDP(R, zeros(2, 3, 2), beta)
        end
    end

    @testset "ddp_negative_inf_error()" begin
        # Dense Matrix
        n, m = 3, 2
        R = [0 1;
             0 -Inf;
            -Inf -Inf]
        Q = fill(1.0/n, n, m, n)
        beta = 0.95

        @test_throws ArgumentError DiscreteDP(R, Q, beta)

        # State-Action Pair Formulation
        #
        # s_indices = [0, 0, 1, 1, 2, 2]
        # a_indices = [0, 1, 0, 1, 0, 1]
        # R_sa = reshape(R, n*m)
        # Q_sa_dense = reshape(Q, n*m, n)          #TODO: @sglyon Not sure how to reshape in Julia
        #
        # @test_throws ArgumentError DiscreteDP(R_sa, Q_sa, beta, s_indices, a_indices)
    end

end # end @testset


text:
# copyright (c) 2017 Ralph Smith
# SPDX-License-Identifier: MIT-Expat
# License-Filename: LICENSE.md

eye(A::Matrix{T}) where T = Matrix{eltype(A)}(I,size(A))

function gnsdcheck(A,B,V,ν,μ,tol=1e-8; quiet=false)
    n = size(B,1)
    @assert size(A) == (n,n)
    @assert size(B) == (n,n)
    @assert size(V) == (n,n)
    @assert ν == length(μ)

    ok = true
    t = norm(A - V * B * V') / norm(A)
    if t > tol
        ok = false
        warn("decomposition error: $t")
    elseif !quiet
        println("decomposition error: ",t)
    end
    t = norm(eye(V) - V' * V)
    if t > 10n*eps()
        ok = false
        warn("unitarity error: $t")
    elseif !quiet
        println("unitarity error: ",t)
    end

    # compute start indices of blocks
    bs = cumsum(vcat([1],μ))
    (bs[end] == n+1) || push!(bs,n+1)
    nb = length(bs)-1

    for j=1:ν
        if norm(B[bs[j]:n, bs[j]:bs[j+1]-1]) != 0
            warn("subdiag block $j is not null")
            ok=false
        end
    end
    for ib=1:nb
        Bii = B[bs[ib]:bs[ib+1]-1,bs[ib]:bs[ib+1]-1]
        if ib>ν
            F = svd!(Bii)
            if any(F.S.< tol)
                warn("diag block $ib (> ν) is singular")
                ok=false
            end
        else
            if norm(Bii) != 0
                warn("nonzero diag block $ib")
                ok=false
            end
        end
        for jb=ib+1:nb
            Bij = B[bs[ib]:bs[ib+1]-1,bs[jb]:bs[jb+1]-1]
            F = svd!(Bij)
            if any(F.S.< tol)
                warn("block $ib,$jb does not have full column rank")
                ok=false
            end
        end
    end
    ok && (quiet || println("all tests pass."))
    ok
end

"""
    makegnb(ν,μx) -> B

construct a matrix with specified generalized nullspace structure

This makes the block-structured part `B`; apply a unitary transform to taste.
"""
function makegnb(ν,μx; cplx=false)
    nb = length(μx)
    (nb ∈ [ν,ν+1]) || throw(ArgumentError("must specify ν or ν+1 blocks"))
    for j=2:ν
        (μx[j-1] >= μx[j]) ||
            throw(ArgumentError("singular block sizes must be nonincreasing"))
    end
    bs = vcat([1],1 .+ cumsum(μx)) # starting indices of blocks
    n = bs[end] -1
    if cplx
        B = randn(n,n) + 1im*randn(n,n)
    else
        B = randn(n,n)
    end
    for i=1:ν
        B[bs[i]:n, bs[i]:bs[i+1]-1] .= 0
    end
    rmul!(B,1/norm(B))
    B
end


text:
# Copyright (c) 2013 Jasper den Ouden, under the MIT license,
# see doc/mit.txt from the project directory.

module Treekenize

import Base.readline

export treekenize, StrExpr #Function for making trees itself.
export none_incorrect
#Each element needs these to know what to do.
export head_expr, head_begin,head_end,head_infix, head_seeker
export infix_str, infix_expr

export infix_it

#Out of place general purpose, probably what i want already exists.
export ConvenientStream

# --- end module stuff

#Makes a stream with a string.
#TODO pull request `stream(str::String)` and `stream(cmd::Cmd)`
function stream_from_string(string::String)
  s = memio(length(string))
  write(s, string)
  seek(s,0) #Need to seek to start.
  return s
end
string_from_stream() = memio() #Which is a IOStream, which has a (todo name it

type ConvenientStream #TODO not quite the right place to put this.
    stream::IOStream
    line::String
    line_n::Int64
end
ConvenientStream(stream::IOStream) = ConvenientStream(stream,"",int64(0))
ConvenientStream(string::String) = ConvenientStream(stream_from_string(string))

#Pass on @with duties.
no_longer_with(cs::ConvenientStream) = no_longer_with(cs.stream)

function forward(cs::ConvenientStream, n::Integer)
#    print("|",cs.line[1:n])
    cs.line = cs.line[n:]
end
function readline(cs::ConvenientStream)
    cs.line_n += 1
    add_line = readline(cs.stream)
    cs.line = string(cs.line, add_line)
end

function first_last_search(in_str, search_str, from)
    range = search(in_str, search_str, from)
    return first(range),last(range)+1
end
first_last_search(in_str, search_str) =
    first_last_search(in_str, search_str,1)

type IncorrectEnd
    initial_n::Int64 #Where it started.
    incorrect_n::Int64 #Where it ended with an incorrect ending symbol.

    correct_end
    got
end

#Infixes a list with given precidence.
function infix_it(list::Array, order::Array)
    if isempty(order)
        return list #length(list)==1 ? list[1] : list
    end
    infix = infix_str(order[1])
    ret = {}
    for el in list
        if isa(el,String)
            i,f,t = 1,1,1
            while true
                f,t = first_last_search(el, infix, i)
                if t==-1 || t==0
                    break
                end
                if i<f
                    push!(ret, infix_it({el[i:f-1]}, order[2:]))
                end
                i = t
            end
            if t!=length(el)
                push!(ret, infix_it({el[i:]}, order[2:]))
            end
        else
            push!(ret, infix_it(el, order[2:]))
        end
    end
    return (length(ret) == 1 ? ret[1] :
                               (isempty(ret) ? {} : infix_expr(order[1], ret)))
end

#Turns a stream into a tree of stuff.
# `stream`  is the input stream
# `which`   is two arrays:
#           elements of the first indicate(extractable with functions-):
#              `head_expr`, produces an 'expression' from a list.
#              `head_begin` what string starts it.
#              `head_end`   what end the expression will have.
#           the second is simply an array of disallowed strings.
# `end_str`      String that ends the current tree.
# `try_cnt`      Max number of attempts reading a line.
# `longest_len`  Longest word that begins/ends.
function treekenize(stream::ConvenientStream, which::(Array,Array),
                    on_head,
                    try_cnt::Integer, longest_len::Integer)
    end_str = head_end(on_head)
    seeker,not_incorrect = which
    list = {}
    n=0
    initial_n = stream.line_n
    readline(stream)

    while n< try_cnt
        pick = nothing
        min_s,min_e = typemax(Int64),0

        search_str = stream.line
        for el in seeker
            s,e = first_last_search(search_str, head_begin(el))
            if s!=0 && s< min_s
                pick = el
                min_s,min_e = s,e
                if min_s + longest_len < length(search_str)
                    search_str =  search_str[1:min_s + longest_len]
                end
            end
        end
        s,e = first_last_search(search_str, end_str)
        search_str = search_str[1:s-1] #Warning about this guy.
        assert(s==0 || s<e)
      #Look for enders that dont match the begin.
      # (Depending on input some may be allowed)
        for el in not_incorrect
            s2,e2 = first_last_search(search_str, el)
            #Shouldnt be inside subtree.
            if s2!=0 && min_e!=0 && s2<min_s
                throw(IncorrectEnd(initial_n, stream.line_n, end_str, el))
            end
        end

        if s!=0 && s<min_s #Ended before some subtree starting symbol.
            n=0
            if s>1
                assert( length(search_str) == s-1 )
                push!(list, search_str) #[1:s-1] (already done)
            end
            forward(stream, e)
            return infix_it(list, head_infix(on_head)) #Go up a level.
        elseif pick==nothing #got nothing, fetch some more.
            n+=1
            readline(stream)
        else #Got new branch.
            n=0
            if min_s>1
                push!(list, stream.line[1:min_s-1]) #Push what is here.
            end
            forward(stream, min_e)
           #Push branch.
            push!(list, head_expr(pick,
                                  treekenize(stream, head_seeker(which,pick),
                                             pick, try_cnt, longest_len)))
        end
    end
    #TODO failed to end everything, this is potentially an error!
    # Problem is that i dont see a good way to check if it is eof.. (no mention)
    return list
end

#Makes the ConvenientStream for you.
treekenize{T}(stream::T, which::(Array,Array), on_head,
           try_cnt::Integer, longest_len::Integer) =
    treekenize(ConvenientStream(stream), which, on_head, try_cnt,longest_len)

#(this version:)not_incorrect defaults to not checking anything.
treekenize{T}(thing::T, seeker::Array, on_head,
              try_cnt::Integer, longest_len::Integer) =
    treekenize(thing, (seeker,{}), on_head, try_cnt,longest_len)

#If _all_ the given seekers may not be present incorrectly, this tries to
# make the list to detect it.
function none_incorrect(seeker::Array)
    list = {}
    for el in seeker
        if isa(el,Tuple) && length(el)>=2 && isa(el[2],String)
            push!(list, el[2])
        end
    end
    return list
end

type StrExpr #Structure for default AST tree.
    head::String
    body::Array
end

#Head is defaultly just the begin.
head_expr{T}(el::T, list::Array) = StrExpr(head_begin(el), list)
head_expr{T}(el::T, first::StrExpr) = head_expr(el, {first})

#The (String,String);(beginner,ender) seeker;
head_begin(el::(String,String))      = el[1]
head_end(el::(String,String))        = el[2]
head_infix(el::(String,String))      = Array(String,0) #No infix notation.
head_seeker(got,el::(String,String)) = got #Just keeps going with the same seeker.

#The (String,String,Array{String,1}); (beginner,ender, infix order)
typealias BEI (ASCIIString,ASCIIString,Array{ASCIIString,1}) #Convenience.
head_begin(el::BEI)      = el[1]
head_end(el::BEI)        = el[2]
head_infix(el::BEI)      = el[3]
head_seeker(got,el::BEI) = got #Just keeps going with the same seeker.

head_begin(x) = error("$x\n$(typeof(x[3]))")

infix_str(str::String) = str
infix_expr(str::String, list::Array) = StrExpr(str, list)

end #module Treekenize


text:
function smolyak_barycentric_points(nodes::Array{T,2},x::Array{T,1},num_points::S) where {T<:AbstractFloat,S<:Integer}

  n = size(nodes,1)

  if num_points > n
    println("smolyak_barycentric_points: Too many points")
    throw(ErrorException())
  end

  dist_vector    = zeros(T,n)
  index_vector_k = Array{S}(num_points)

  for i = 1:n
    dist_vector[i] = norm(x-nodes[i,:])
  end

  for i = 1:num_points
    (dist,index)       = findmin(dist_vector)
    index_vector_k[i]  = index
    dist_vector[index] = Inf
  end

  return index_vector_k

end

function smolyak_barycentric_evaluate(y::AbstractArray{T,1},nodes::Array{T,2},x::Array{T,1}) where T<:AbstractFloat

  k = size(nodes,2)+1
  index_vector_k = smolyak_barycentric_points(nodes,x,k)

  l = vcat(nodes[index_vector_k,:]',ones(1,k))
  println(l)
  weights = l\[x;1.0]

  y_hat = zero(T)

  for i = 1:k
    y_hat += y[index_vector_k[i]]*weights[i]
  end

  return y_hat

end


text:
orth(S::AbstractArray{<:Real}) = Matrix(qr(S).Q)

"""
    S_update_a

Random set of orthonormal sample directions.

See: Equation (6.1.a).
"""
function S_update_a(Sₖ, Yₖ, pₖ)
    return orth(randn(eltype(Sₖ), size(Sₖ)...))
end

"""
    S_update_b

Random set of sample directions orthogonal to the previous sample space given by input Sₖ.

See: Equation (6.1.b).
"""
function S_update_b(Sₖ, Yₖ, pₖ)
    M = randn(eltype(Sₖ), size(Sₖ)...)
    return orth(M - Sₖ * (Sₖ' * M))
end

"""
    S_update_c

Attempts to guide the algorithm in accurately resolving the eigenspace associated with the
larger Hessian eigenvalues.

See: Equation (6.1.c).
"""
function S_update_c(Sₖ, Yₖ, pₖ)
    return orth(Yₖ - Sₖ * (Sₖ' * Yₖ))
end

"""
    S_update_d

A variant of (6.1.a) that includes approximate curvature information along the previously
chosen step.

See: Equation (6.1.d).
"""
function S_update_d(Sₖ, Yₖ, pₖ)
    return orth([orth(randn(eltype(Sₖ), size(Sₖ, 1), size(Sₖ, 2) - 1)) pₖ])
end

"""
    S_update_e

A variant of (6.1.b) that includes approximate curvature information along the previously
chosen step.

See: Equation (6.1.e).
"""
function S_update_e(Sₖ, Yₖ, pₖ)
    M = randn(size(Sₖ, 1), size(Sₖ, 2) - 1)
    return orth([orth(M - Sₖ * (Sₖ' * M)) pₖ])
end

"""
    S_update_f

A variant of (6.1c) that includes approximate curvature information along the previously
chosen step.

See: Equation (6.1.f).
"""
function S_update_f(Sₖ, Yₖ, pₖ)
    return orth([orth(Yₖ[:, begin:end-1] - Sₖ * (Sₖ' * Yₖ[:, begin:end-1])) pₖ])
end

"""
    SR1

Returns the algebraically minimal SR1 inverse Quasi-Newton block update,
where ``δ`` is the Moore-Penrose pseudoinverse relative tolerance used in `pinv`.

See: Algorithm ``4.2.``
"""
function SR1(
    H::AbstractArray{<:Real},
    U::AbstractArray{<:Real},
    V::AbstractArray{<:Real},
    δ::Float64,
)
    U_minus_HV = U - H * V

    if size(U, 2) == 1
        return Symmetric(H + ((U_minus_HV) * (U_minus_HV)') / ((U_minus_HV)' * V))
    end

    return Symmetric(H + U_minus_HV * pinv(U_minus_HV' * V, rtol = δ) * U_minus_HV')
end

"""
    PSB

Powell-Symmetric-Broyden generalized Quasi-Newton block update,
where ``δ`` is the Moore-Penrose pseudoinverse relative tolerance used in `pinv`.

See: Algorithm ``4.3.``
"""
function PSB(
    H::AbstractArray{<:Real},
    U::AbstractArray{<:Real},
    V::AbstractArray{<:Real},
    δ::Float64,
)
    if size(V, 2) == 1
        T₁ = 1 / (V' * V)
    else
        T₁ = pinv(V' * V, rtol = δ)
    end

    T₂ = V * T₁ * (U - H * V)'
    return Symmetric(H + T₂ + T₂' - T₂ * V * T₁ * V')
end

"""
Driver

An immutable structure specifying a simulation's preliminary secant update flag (`pflag`),
the ``Sₖ₊₁`` update formula, and the Quasi-Newton update formula used. If keyword arguments
aren't provided to a drivers' construction, then the default `S_update`, `pflag`, and `QN_updates`
are assigned. A Driver contains a mutable `DriverOptions` instance, which allows for forwarding
of the options interface to the Driver interface.

See `S_update`, `QN_update`, and `pflag` for more on keyword arguments.
"""
mutable struct Driver
    S_update::Function
    QN_update::Function
    pflag::Bool
    options::DriverOptions

    function Driver(;
        S_update = S_update_c,
        QN_update = SR1,
        pflag = false,
        options = DriverOptions(),
    )
        # TODO: Confrim S_update and QN_update inputs are valid
        #       Maybe parse symbol and compare to set of strings?
        new(S_update, QN_update, pflag, options)
    end
end

"""
    S_update(d::Driver)

The supplemental sample direction update formula of Driver `d`.

Values: `S_update_a`, `S_update_b`, `S_update_c`, `S_update_d`, `S_update_e`, `S_update_f`
"""
S_update(d::Driver) = getfield(d, :S_update)

"""
    QN_update(d::Driver)

The QN update formula of Driver `d`.

Values: `SR1`, `PSB`
"""
QN_update(d::Driver) = getfield(d, :QN_update)

"""
    pflag(d::Driver)

The preliminary secant QN update flag of driver `d` wraping a boolean.

Options `true`, `false`
"""
pflag(d::Driver) = getfield(d, :pflag)

options(d::Driver) = getfield(d, :options)

# Inhereted/Forwarded DriverOptions Methods
samples(d::Driver) = samples(options(d))

samples!(d::Driver, samples) = samples!(options(d), samples)

Δ_max(d::Driver) = Δ_max(options(d))

Δ_max!(d::Driver, Δ_max) = Δ_max!(options(d), Δ_max)

δ_tol(d::Driver) = δ_tol(options(d))

δ_tol!(d::Driver, δ_tol) = δ_tol!(options(d), δ_tol)

ϵ_tol(d::Driver) = ϵ_tol(options(d))

ϵ_tol!(d::Driver, ϵ_tol) = ϵ_tol!(options(d), ϵ_tol)

max_iterations(d::Driver) = max_iterations(options(d))

max_iterations!(d::Driver, m) = max_iterations!(options(d), m)

weave_level(d::Driver) = weave_level(options(d))

weave_level!(d::Driver, level) = weave_level!(options(d), level)

log_level(d::Driver) = log_level(options(d))

log_level!(d::Driver, level) = log_level!(options(d), level)

Base.getproperty(d::Driver, s::Symbol) = @restrict Driver

Base.propertynames(d::Driver) = ()


text:
using DifferentialEquations
using Plots
using DataFrames
using CSV

abstract type DynamicalSystem end

struct ODE1 <: DynamicalSystem end
struct ODE3 <: DynamicalSystem end
struct LorenzSystem <: DynamicalSystem end
struct LotkaVolterra <: DynamicalSystem end
struct NSECylinder <: DynamicalSystem end
struct HelicopterData <: DynamicalSystem
    filename::String
    HelicopterData(file) = new(file)
end

function noise(u::Vector{Vector{T}}) where T
    for i = 1:length(u)
        for j = 1:length(u[i])
            u[i][j] += 0.1*randn()
        end
    end
    return u
end

function datagen(ds::ODE1)
    function ode1!(du,u,p,t)
        du[1] = -0.1*u[1] + 2*u[2]
        du[2] = -2*u[2] - 0.1*u[2]
    end

    u₀ = [1,1]
    p = [0.7,0.3,0.3,0.4]
    tspan = (0,20)
    t = 0.01:0.01:20
    prob = ODEProblem(ode1!,u₀,tspan,p)
    sol = solve(prob,Tsit5(),saveat=t)
    # noisy = noise(sol.u)
    fig = plot(sol,vars=(1,2),title="ODE System Linear")
    display(fig)
    return sol
end

function datagen(ds::ODE3)
    function ode1!(du,u,p,t)
        du[1] = -0.1*u[1]^3 + 2*u[2]^3
        du[2] = -2*u[2]^3 - 0.1*u[2]^3
    end

    u₀ = [1,1]
    p = [0.7,0.3,0.3,0.4]
    tspan = (0,20)
    t = 0.01:0.01:20
    prob = ODEProblem(ode1!,u₀,tspan,p)
    sol = solve(prob,Tsit5(),saveat=t)
    # noisy = noise(sol.u)
    fig = plot(sol,vars=(1,2),title="ODE System Non Linear")
    display(fig)
    return sol
end

function datagen(data::LorenzSystem)
    function lorenz!(du,u,p,t)
        du[1] = p[1]*(u[2]-u[1])
        du[2] = u[1]*(p[2]-u[3]) - u[2]
        du[3] = u[1]*u[2] - p[3]*u[3]
    end

    u₀ = [1,0,0]
    p = [10,28,8/3] # Change d_lorenz
    tspan = (0,100)
    t = 0.01:0.01:100
    prob = ODEProblem(lorenz!,u₀,tspan,p)
    sol = solve(prob,Tsit5(),saveat=t)
    # noisy = noise(sol.u)
    fig = plot(sol,vars=(1,2,3),title="Lorenz Attractor")
    display(fig)
    return sol
end

function datagen(ds::LotkaVolterra)
    function lv!(du,u,p,t)
        α,β,γ,δ = p
        du[1] = α*u[1] - β*u[1]*u[2]
        du[2] = - γ*u[2] + δ*u[1]*u[2]
    end

    u₀ = [1,1]
    p = [0.7,0.3,0.3,0.4]
    tspan = (0,20)
    t = 0.01:0.01:20
    prob = ODEProblem(lv!,u₀,tspan,p)
    sol = solve(prob,Tsit5(),saveat=t)
    # noisy = noise(sol.u)
    fig = plot(sol,vars=(1,2),title="Lotka Volterra")
    display(fig)
    return sol
end

function datagen(ds::HelicopterData)
    df = CSV.read(ds.filename,DataFrame,delim=",")
    data = Matrix(df)
    data = [data[:,3] data[:,4]]
    my_cgrad = cgrad([:blue, :green, :red])
    x = range(0,stop=131,length=size(data,1))
    fig = plot(data[:,1],data[:,2],lc=my_cgrad,line_z=x,title="Helicopter data")
    display(fig)
    return data
end


text:
using ParticleGridMapper
using Documenter

DocMeta.setdocmeta!(ParticleGridMapper, :DocTestSetup, :(using ParticleGridMapper); recursive=true)

makedocs(;
    modules=[ParticleGridMapper],
    authors="Chia-Yu Hu <cyhu.astro@gmail.com> and contributors",
    repo="https://github.com/huchiayu/ParticleGridMapper.jl/blob/{commit}{path}#{line}",
    sitename="ParticleGridMapper.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://huchiayu.github.io/ParticleGridMapper.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
    ],
)

deploydocs(;
    repo="github.com/huchiayu/ParticleGridMapper.jl",
)


text:
#######################################

immutable LaneTag
    segment :: Int # segment id
    lane    :: Int # index in segment.lanes of this lane
end
Base.show(io::IO, tag::LaneTag) = @printf(io, "LaneTag(%d, %d)", tag.segment, tag.lane)
hash(id::LaneTag, h::UInt=zero(UInt)) = hash(id.segment, hash(id.lane, h))
@compat Base.:(==)(a::LaneTag, b::LaneTag) = a.segment == b.segment && a.lane == b.lane
const NULL_LANETAG = LaneTag(0,0)

#######################################

immutable RoadIndex
    ind::CurveIndex
    tag::LaneTag
end
const NULL_ROADINDEX = RoadIndex(CurveIndex(-1,NaN), LaneTag(-1,-1))

Base.show(io::IO, r::RoadIndex) = @printf(io, "RoadIndex({%d, %3.f}, {%d, %d})", r.ind.i, r.ind.t, r.tag.segment, r.tag.lane)
Base.write(io::IO, r::RoadIndex) = @printf(io, "%d %.6f %d %d", r.ind.i, r.ind.t, r.tag.segment, r.tag.lane)

#######################################

immutable LaneConnection
    downstream::Bool # if true, mylane → target, else target → mylane
    mylane::CurveIndex
    target::RoadIndex
end

Base.show(io::IO, c::LaneConnection) = print(io, "LaneConnection(", c.downstream ? "D" : "U", ", ", c.mylane, ", ", c.target)
function Base.write(io::IO, c::LaneConnection)
    @printf(io, "%s (%d %.6f) ", c.downstream ? "D" : "U", c.mylane.i, c.mylane.t)
    write(io, c.target)
end
function Base.parse(::Type{LaneConnection}, line::AbstractString)
    cleanedline = replace(line, r"(\(|\))", "")
    tokens = split(cleanedline, ' ')

    @assert(tokens[1] == "D" || tokens[1] == "U")
    downstream = tokens[1] == "D"
    mylane = CurveIndex(parse(Int, tokens[2]), parse(Float64, tokens[3]))
    target = RoadIndex(
                CurveIndex(parse(Int, tokens[4]), parse(Float64, tokens[5])),
                LaneTag(parse(Int, tokens[6]), parse(Int, tokens[7]))
            )
    LaneConnection(downstream, mylane, target)
end

#######################################

const DEFAULT_LANE_WIDTH = 3.0 # [m]

type Lane
    tag            :: LaneTag
    curve          :: Curve
    width          :: Float64 # [m]
    speed_limit    :: SpeedLimit
    boundary_left  :: LaneBoundary
    boundary_right :: LaneBoundary
    exits          :: Vector{LaneConnection} # list of exits; put the primary exit (at end of lane) first
    entrances      :: Vector{LaneConnection} # list of entrances; put the primary entrance (at start of lane) first

    function Lane(
        tag::LaneTag,
        curve::Curve;
        width::Float64 = DEFAULT_LANE_WIDTH,
        speed_limit::SpeedLimit = DEFAULT_SPEED_LIMIT,
        boundary_left::LaneBoundary = NULL_BOUNDARY,
        boundary_right::LaneBoundary = NULL_BOUNDARY,
        exits::Vector{LaneConnection} = LaneConnection[],
        entrances::Vector{LaneConnection} = LaneConnection[],
        next::RoadIndex=NULL_ROADINDEX,
        prev::RoadIndex=NULL_ROADINDEX,
        )

        retval = new()
        retval.tag = tag
        retval.curve = curve
        retval.width = width
        retval.speed_limit = speed_limit
        retval.boundary_left = boundary_left
        retval.boundary_right = boundary_right
        retval.exits = exits
        retval.entrances = entrances

        if next != NULL_ROADINDEX
            unshift!(retval.exits, LaneConnection(true, curveindex_end(retval.curve), next))
        end
        if prev != NULL_ROADINDEX
            unshift!(retval.entrances, LaneConnection(false, CURVEINDEX_START, prev))
        end

        retval
    end
end

has_next(lane::Lane) = !isempty(lane.exits) && lane.exits[1].mylane == curveindex_end(lane.curve)
has_prev(lane::Lane) = !isempty(lane.entrances) && lane.entrances[1].mylane == CURVEINDEX_START

is_in_exits(lane::Lane, target::LaneTag) = findfirst(lc->lc.target.tag == target, lane.exits) != 0
is_in_entrances(lane::Lane, target::LaneTag) = findfirst(lc->lc.target.tag == target, lane.entrances) != 0

function connect!(source::Lane, dest::Lane)
    # place these at the front

    cindS = curveindex_end(source.curve)
    cindD = CURVEINDEX_START

    unshift!(source.exits,   LaneConnection(true,  cindS, RoadIndex(cindD, dest.tag)))
    unshift!(dest.entrances, LaneConnection(false, cindD, RoadIndex(cindS, source.tag)))
    (source, dest)
end
function connect!(source::Lane, cindS::CurveIndex, dest::Lane, cindD::CurveIndex)
    push!(source.exits,   LaneConnection(true,  cindS, RoadIndex(cindD, dest.tag)))
    push!(dest.entrances, LaneConnection(false, cindD, RoadIndex(cindS, source.tag)))
    (source, dest)
end

#######################################

type RoadSegment
    # a list of lanes forming a single road with a common direction
    id::Int
    lanes::Vector{Lane} # lanes are stored right to left
    RoadSegment(id::Int, lanes::Vector{Lane}=Lane[]) = new(id, lanes)
end

#######################################

type Roadway
    segments::Vector{RoadSegment}
    Roadway(segments::Vector{RoadSegment}=RoadSegment[]) = new(segments)
end

function Base.write(io::IO, roadway::Roadway)
    # writes to a text file

    println(io, "ROADWAY")
    println(io, length(roadway.segments)) # number of segments
    for seg in roadway.segments
        println(io, seg.id)
        println(io, "\t", length(seg.lanes)) # number of lanes
        for (i,lane) in enumerate(seg.lanes)
            @assert(lane.tag.lane == i)
            @printf(io, "\t%d\n", i)
            @printf(io, "\t\t%.3f\n", lane.width)
            @printf(io, "\t\t%.3f %.3f\n", lane.speed_limit.lo, lane.speed_limit.hi)
            println(io, "\t\t", lane.boundary_left.style, " ", lane.boundary_left.color)
            println(io, "\t\t", lane.boundary_right.style, " ", lane.boundary_right.color)
            println(io, "\t\t", length(lane.exits) + length(lane.entrances))
            for conn in lane.exits
                print(io, "\t\t\t"); write(io, conn); print(io, "\n")
            end
            for conn in lane.entrances
                print(io, "\t\t\t"); write(io, conn); print(io, "\n")
            end
            println(io, "\t\t", length(lane.curve))
            for pt in lane.curve
                @printf(io, "\t\t\t(%.4f %.4f %.6f) %.4f %.8f %.8f\n", pt.pos.x, pt.pos.y, pt.pos.θ, pt.s, pt.k, pt.kd)
            end
        end
    end
end
function Base.read(io::IO, ::Type{Roadway})
    lines = readlines(io)
    line_index = 1
    if contains(lines[line_index], "ROADWAY")
        line_index += 1
    end

    function advance!()
        line = strip(lines[line_index])
        line_index += 1
        line
    end

    nsegs = parse(Int, advance!())
    roadway = Roadway(Array{RoadSegment}(nsegs))
    for i_seg in 1:nsegs
        segid = parse(Int, advance!())
        nlanes = parse(Int, advance!())
        seg = RoadSegment(segid, Array{Lane}(nlanes))
        for i_lane in 1:nlanes
            @assert(i_lane == parse(Int, advance!()))
            tag = LaneTag(segid, i_lane)
            width = parse(Float64, advance!())

            tokens = split(advance!(), ' ')
            speed_limit = SpeedLimit(parse(Float64, tokens[1]), parse(Float64, tokens[2]))

            tokens = split(advance!(), ' ')
            boundary_left = LaneBoundary(Symbol(tokens[1]), Symbol(tokens[2]))

            tokens = split(advance!(), ' ')
            boundary_right = LaneBoundary(Symbol(tokens[1]), Symbol(tokens[2]))

            exits = LaneConnection[]
            entrances = LaneConnection[]
            n_conns = parse(Int, advance!())
            for i_conn in 1:n_conns
                conn = parse(LaneConnection, advance!())
                conn.downstream ? push!(exits, conn) : push!(entrances, conn)
            end

            npts = parse(Int, advance!())
            curve = Array{CurvePt}(npts)
            for i_pt in 1:npts
                line = advance!()
                cleanedline = replace(line, r"(\(|\))", "")
                tokens = split(cleanedline, ' ')
                x = parse(Float64, tokens[1])
                y = parse(Float64, tokens[2])
                θ = parse(Float64, tokens[3])
                s = parse(Float64, tokens[4])
                k = parse(Float64, tokens[5])
                kd = parse(Float64, tokens[6])
                curve[i_pt] = CurvePt(VecSE2(x,y,θ), s, k, kd)
            end

            seg.lanes[i_lane] = Lane(tag, curve, width=width, speed_limit=speed_limit,
                                     boundary_left=boundary_left,
                                     boundary_right=boundary_right,
                                     entrances=entrances, exits=exits)
        end
        roadway.segments[i_seg] = seg
    end

    roadway
end

"""
    lane[ind::CurveIndex, roadway::Roadway]
Accessor for lanes based on a CurveIndex.
Note that we extend the definition of a CurveIndex,
previously ind.i ∈ [1, length(curve)-1], to:

    ind.i ∈ [0, length(curve)]

where 1 ≤ ind.i ≤ length(curve)-1 is as before, but if the index
is on the section between two lanes, we use:

    ind.i = length(curve), ind.t ∈ [0,1] for the region between curve[end] → next
    ind.i = 0,             ind.t ∈ [0,1] for the region between prev → curve[1]
"""
function Base.getindex(lane::Lane, ind::CurveIndex, roadway::Roadway)
    if ind.i == 0
        pt_lo = prev_lane_point(lane, roadway)
        pt_hi = lane.curve[1]
        s_gap = abs(pt_hi.pos - pt_lo.pos)
        pt_lo = CurvePt(pt_lo.pos, -s_gap, pt_lo.k, pt_lo.kd)
        lerp(pt_lo, pt_hi, ind.t)
    elseif ind.i < length(lane.curve)
        lane.curve[ind]
    else
        pt_hi = next_lane_point(lane, roadway)
        pt_lo = lane.curve[end]
        s_gap = abs(pt_hi.pos - pt_lo.pos)
        pt_hi = CurvePt(pt_hi.pos, pt_lo.s + s_gap, pt_hi.k, pt_hi.kd)
        lerp( pt_lo, pt_hi, ind.t)
    end
end
function Base.getindex(roadway::Roadway, segid::Int)
    for seg in roadway.segments
        if seg.id == segid
            return seg
        end
    end
    error("Could not find segid $segid in roadway")
end
function Base.getindex(roadway::Roadway, tag::LaneTag)
    seg = roadway[tag.segment]
    seg.lanes[tag.lane]
end

is_between_segments_lo(ind::CurveIndex) = ind.i == 0
is_between_segments_hi(ind::CurveIndex, curve::Curve) = ind.i == length(curve)
is_between_segments(ind::CurveIndex, curve::Curve) = is_between_segments_lo(ind) || is_between_segments_hi(ind, curve)

next_lane(lane::Lane, roadway::Roadway) = roadway[lane.exits[1].target.tag]
prev_lane(lane::Lane, roadway::Roadway) = roadway[lane.entrances[1].target.tag]
next_lane_point(lane::Lane, roadway::Roadway) = roadway[lane.exits[1].target]
prev_lane_point(lane::Lane, roadway::Roadway) = roadway[lane.entrances[1].target]

function has_segment(roadway::Roadway, segid::Int)
    for seg in roadway.segments
        if seg.id == segid
            return true
        end
    end
    false
end
function has_lanetag(roadway::Roadway, tag::LaneTag)
    if !has_segment(roadway, tag.segment)
        return false
    end
    seg = roadway[tag.segment]
    1 ≤ tag.lane ≤ length(seg.lanes)
end

immutable RoadProjection
    curveproj::CurveProjection
    tag::LaneTag
end

function get_closest_perpendicular_point_between_points(A::VecSE2, B::VecSE2, Q::VecSE2;
    tolerance::Float64 = 0.01, # acceptable error in perpendicular component
    max_iter::Int = 50, # maximum number of iterations
    )

    # CONDITIONS: a < b, either f(a) < 0 and f(b) > 0 or f(a) > 0 and f(b) < 0
    # OUTPUT: value which differs from a root of f(x)=0 by less than TOL

    a = 0.0
    b = 1.0

    f_a = inertial2body(Q, A).x
    f_b = inertial2body(Q, B).x

    if sign(f_a) == sign(f_b) # both are wrong - use the old way
        t = get_lerp_time_unclamped(A, B, Q)
        t = clamp(t, 0.0, 1.0)
        return (t, lerp(A,B,t))
    end

    iter = 1
    while iter ≤ max_iter
        c = (a+b)/2 # new midpoint
        footpoint = lerp(A, B, c)
        f_c = inertial2body(Q, footpoint).x
        if abs(f_c) < tolerance # solution found
            return (c, footpoint)
        end
        if sign(f_c) == sign(f_a)
            a, f_a = c, f_c
        else
            b = c
        end
        iter += 1
    end

    # Maximum number of iterations passed
    # This will occur when we project with a point that is not actually in the range,
    # and we converge towards one edge

    if a == 0.0
        return (0.0, A)
    elseif b == 1.0
        return (1.0, B)
    else
        warn("get_closest_perpendicular_point_between_points - should not happen")
        c = (a+b)/2 # should not happen
        return (c, lerp(A,B,c))
    end
end

"""
    proj(posG::VecSE2, lane::Lane, roadway::Roadway)
Return the RoadProjection for projecting posG onto the lane.
This will automatically project to the next or prev curve as appropriate.
"""
function Vec.proj(posG::VecSE2, lane::Lane, roadway::Roadway;
    move_along_curves::Bool = true, # if false, will only project to lane.curve
    )
    curveproj = proj(posG, lane.curve)
    rettag = lane.tag

    if curveproj.ind == CurveIndex(1,0.0) && has_prev(lane)
        pt_lo = prev_lane_point(lane, roadway)
        pt_hi = lane.curve[1]

        t = get_lerp_time_unclamped(pt_lo, pt_hi, posG)
        if t ≤ 0.0 && move_along_curves
            return proj(posG, prev_lane(lane, roadway), roadway)
        elseif t < 1.0 # for t == 1.0 we use the actual end of the lane
            @assert(!move_along_curves || 0.0 ≤ t < 1.0)

            # t was computed assuming a constant angle
            # this is not valid for the large distances and angle disparities between lanes
            # thus we now use a bisection search to find the appropriate location

            t, footpoint = get_closest_perpendicular_point_between_points(pt_lo.pos, pt_hi.pos, posG)

            ind = CurveIndex(0, t)
            curveproj = get_curve_projection(posG, footpoint, ind)
        end
    elseif curveproj.ind == curveindex_end(lane.curve) && has_next(lane)
        pt_lo = lane.curve[end]
        pt_hi = next_lane_point(lane, roadway)

        t = get_lerp_time_unclamped(pt_lo, pt_hi, posG)

        if t ≥ 1.0 && move_along_curves
             # for t == 1.0 we use the actual start of the lane
            return proj(posG, next_lane(lane, roadway), roadway)
        elseif t ≥ 0.0
            @assert(!move_along_curves || 0.0 ≤ t ≤ 1.0)

            # t was computed assuming a constant angle
            # this is not valid for the large distances and angle disparities between lanes
            # thus we now use a bisection search to find the appropriate location

            t, footpoint = get_closest_perpendicular_point_between_points(pt_lo.pos, pt_hi.pos, posG)

            ind = CurveIndex(length(lane.curve), t)
            curveproj = get_curve_projection(posG, footpoint, ind)
        end
    end
    RoadProjection(curveproj, rettag)
end

"""
    proj(posG::VecSE2, seg::RoadSegment, roadway::Roadway)
Return the RoadProjection for projecting posG onto the segment.
Tries all of the lanes and gets the closest one
"""
function Vec.proj(posG::VecSE2, seg::RoadSegment, roadway::Roadway)

    best_dist2 = Inf
    best_proj = RoadProjection(CurveProjection(CurveIndex(-1,-1), NaN, NaN), NULL_LANETAG)

    for lane in seg.lanes
        roadproj = proj(posG, lane, roadway)
        footpoint = roadway[roadproj.tag][roadproj.curveproj.ind, roadway]
        dist2 = abs2(posG - footpoint.pos)
        if dist2 < best_dist2
            best_dist2 = dist2
            best_proj = roadproj
        end
    end

    best_proj
end

"""
    proj(posG::VecSE2, seg::RoadSegment, roadway::Roadway)
Return the RoadProjection for projecting posG onto the roadway.
Tries all of the lanes and gets the closest one
"""
function Vec.proj(posG::VecSE2, roadway::Roadway)

    best_dist2 = Inf
    best_proj = RoadProjection(CurveProjection(CurveIndex(-1,-1), NaN, NaN), NULL_LANETAG)

    for seg in roadway.segments
        for lane in seg.lanes
            roadproj = proj(posG, lane, roadway, move_along_curves=false)
            targetlane = roadway[roadproj.tag]
            footpoint = targetlane[roadproj.curveproj.ind, roadway]
            dist2 = abs2(posG - footpoint.pos)
            if dist2 < best_dist2
                best_dist2 = dist2
                best_proj = roadproj
            end
        end
    end

    best_proj
end

############################################

RoadIndex(roadproj::RoadProjection) = RoadIndex(roadproj.curveproj.ind, roadproj.tag)

function Base.getindex(roadway::Roadway, roadind::RoadIndex)
    lane = roadway[roadind.tag]
    lane[roadind.ind, roadway]
end

"""
    move_along(roadind::RoadIndex, road::Roadway, Δs::Float64)
Return the RoadIndex at ind's s position + Δs
"""
function move_along(roadind::RoadIndex, roadway::Roadway, Δs::Float64, depth::Int=0)

    lane = roadway[roadind.tag]
    curvept = lane[roadind.ind, roadway]

    if curvept.s + Δs < 0.0
        if has_prev(lane)
            pt_lo = prev_lane_point(lane, roadway)
            pt_hi = lane.curve[1]
            s_gap = abs(pt_hi.pos - pt_lo.pos)

            if curvept.s + Δs < -s_gap
                lane_prev = prev_lane(lane, roadway)
                curveind = curveindex_end(lane_prev.curve)
                roadind = RoadIndex(curveind, lane_prev.tag)
                return move_along(roadind, roadway, Δs + curvept.s + s_gap, depth+1)
            else # in the gap between lanes
                t = (s_gap + curvept.s + Δs) / s_gap
                curveind = CurveIndex(0, t)
                RoadIndex(curveind, lane.tag)
            end

        else # no prev lane, return the beginning of this one
            curveind = CurveIndex(1, 0.0)
            return RoadIndex(curveind, roadind.tag)
        end
    elseif curvept.s + Δs > lane.curve[end].s
        if has_next(lane)
            pt_lo = lane.curve[end]
            pt_hi = next_lane_point(lane, roadway)
            s_gap = abs(pt_hi.pos - pt_lo.pos)

            if curvept.s + Δs ≥ pt_lo.s + s_gap # extends beyond the gap
                curveind = lane.exits[1].target.ind
                roadind = RoadIndex(curveind, lane.exits[1].target.tag)
                return move_along(roadind, roadway, Δs - (lane.curve[end].s + s_gap - curvept.s))
            else # in the gap between lanes
                t = (Δs - (lane.curve[end].s - curvept.s)) / s_gap
                curveind = CurveIndex(0, t)
                RoadIndex(curveind, lane.exits[1].target.tag)
            end
        else # no next lane, return the end of this lane
            curveind = curveindex_end(lane.curve)
            return RoadIndex(curveind, roadind.tag)
        end
    else
        if roadind.ind.i == 0
            ind = get_curve_index(CurveIndex(1,0.0), lane.curve, curvept.s+Δs)
        elseif roadind.ind.i == length(lane.curve)
            ind = get_curve_index(curveindex_end(lane.curve), lane.curve, curvept.s+Δs)
        else
            ind = get_curve_index(roadind.ind, lane.curve, Δs)
        end
        RoadIndex(ind, roadind.tag)
    end
end

n_lanes_right(lane::Lane, roadway::Roadway) = lane.tag.lane - 1
function n_lanes_left(lane::Lane, roadway::Roadway)
    seg = roadway[lane.tag.segment]
    length(seg.lanes) - lane.tag.lane
end

############################################

function _fit_curve(
    pts::AbstractMatrix{Float64}, # 2×n
    desired_distance_between_samples::Real;
    max_iterations::Int=50,
    epsilon::Float64=1e-4,
    n_intervals_in_arclen::Int=100,
    )

    @assert(size(pts, 1) == 2)

    spline_coeffs = fit_cubic_spline(pts)

    L = calc_curve_length(spline_coeffs[1], spline_coeffs[2], n_intervals_per_segment=n_intervals_in_arclen)
    n = round(Int, L/desired_distance_between_samples, RoundNearestTiesUp)+1

    s_arr = collect(linspace(0.0,L,n))
    t_arr = calc_curve_param_given_arclen(spline_coeffs[1], spline_coeffs[2], s_arr,
        curve_length=L, max_iterations=max_iterations, epsilon=epsilon, n_intervals_in_arclen=n_intervals_in_arclen)

    x_arr = sample_spline(spline_coeffs[1], t_arr)
    y_arr = sample_spline(spline_coeffs[2], t_arr)
    θ_arr = sample_spline_theta(spline_coeffs[1], spline_coeffs[2], t_arr)

    κ_arr = sample_spline_curvature(spline_coeffs[1], spline_coeffs[2], t_arr)
    κd_arr = sample_spline_derivative_of_curvature(spline_coeffs[1], spline_coeffs[2], t_arr)

    @assert(!any(s->isnan(s), s_arr))
    @assert(!any(s->isnan(s), x_arr))
    @assert(!any(s->isnan(s), y_arr))
    @assert(!any(s->isnan(s), θ_arr))

    curve = Array{CurvePt}(n)
    for i in 1 : n
        pos = VecSE2(x_arr[i], y_arr[i], θ_arr[i])
        curve[i] = CurvePt(pos, s_arr[i], κ_arr[i], κd_arr[i])
    end
    curve
end

"""
    read_dxf(io::IO, ::Type{Roadway})
Return a Roadway generated from a DXF file

    Layers with names such as seg001 will contain LWPOLYLINEs.
    Each LWPOLYLINE corresponds to a lane centerline, which together
    are all neighbored.
"""
function read_dxf(io::IO, ::Type{Roadway};
    dist_threshold_lane_connect::Float64 = 0.25, # [m]
    desired_distance_between_curve_samples::Float64 = 1.0; # [m]
    )

    lines = readlines(io)

    i = findfirst(lines, "ENTITIES\n")
    i != 0 || error("ENTITIES section not found")

    ###################################################
    # Pull pts for each lane
    lane_pts_dict = Dict{LaneTag, Vector{VecE2}}()

    i = findnext(lines, "LWPOLYLINE\n", i)
    while i != 0
        i = findnext(lines, "  8\n", i)
        if i != 0 # segment identifier found in LWPOLYLINE

            if ismatch(r"(?<=seg)(\d*)", lines[i+1])
                segid = parse(Int, match(r"(?<=seg)(\d*)", lines[i+1]).match)

                    i = findnext(lines, "AcDbPolyline\n", i)
                i != 0 || error("AcDbPolyline not found in LWPOLYLINE!")
                i = findnext(lines, " 90\n", i)
                i != 0 || error("Number of vertices not found in AcDbPolyline!")

                N = parse(Int, lines[i+1])
                N > 0 || error("Empty line segment!")

                pts = Array{VecE2}(N)

                i = findnext(lines, " 10\n", i)
                i != 0 || error("First point not found in AcDbPolyline!")

                for j in 1 : N
                    x = parse(Float64, lines[i+1])
                    y = parse(Float64, lines[i+3])
                    i += 4
                    pts[j] = VecE2(x,y)
                end

                laneid = 1
                for tag in keys(lane_pts_dict)
                    if tag.segment == segid
                        laneid += 1
                    end
                end
                lane_pts_dict[LaneTag(segid, laneid)] = pts
            end

            i = findnext(lines, "LWPOLYLINE\n", i)
        end
    end

    ###################################################
    # Shift pts to connect to previous / next pts
    lane_next_dict = Dict{LaneTag, Tuple{VecE2, LaneTag}}()
    lane_prev_dict = Dict{LaneTag, Tuple{VecE2, LaneTag}}()

    for (tag, pts) in lane_pts_dict
        # see if can connect to next
        best_tag = NULL_LANETAG
        best_ind = -1
        best_sq_dist = dist_threshold_lane_connect
        for (tag2, pts2) in lane_pts_dict
            if tag2.segment != tag.segment
                for (ind,pt) in enumerate(pts2)
                    sq_dist = abs2(pt - pts[end])
                    if sq_dist < best_sq_dist
                        best_sq_dist = sq_dist
                        best_ind = ind
                        best_tag = tag2
                    end
                end
            end
        end
        if best_tag != NULL_LANETAG
            # remove our last pt and set next to pt to their pt
            pop!(pts)
            lane_next_dict[tag] = (lane_pts_dict[best_tag][best_ind], best_tag)
            if best_ind == 1 # set connect prev as well
                lane_prev_dict[best_tag] = (pts[end], tag)
            end
        end
    end
    for (tag, pts) in lane_pts_dict
        # see if can connect to prev
        if !haskey(lane_prev_dict, tag)
            best_tag = NULL_LANETAG
            best_ind = -1
            best_sq_dist = dist_threshold_lane_connect
            for (tag2, pts2) in lane_pts_dict
                if tag2.segment != tag.segment
                    for (ind,pt) in enumerate(pts2)
                        sq_dist = abs2(pt - pts[1])
                        if sq_dist < best_sq_dist
                            best_sq_dist = sq_dist
                            best_ind = ind
                            best_tag = tag2
                        end
                    end
                end
            end
            if best_tag != NULL_LANETAG
                # connect 'em
                shift!(pts)
                lane_prev_dict[tag] = (lane_pts_dict[best_tag][best_ind], best_tag)
            end
        end
    end

    ###################################################
    # Build the roadway
    retval = Roadway()
    for (tag, pts) in lane_pts_dict
        if !has_segment(retval, tag.segment)
            push!(retval.segments, RoadSegment(tag.segment))
        end
    end

    lane_new_dict = Dict{LaneTag, LaneTag}() # old -> new tag
    for seg in retval.segments

        # pull lanetags for this seg
        lanetags = LaneTag[]
        for tag in keys(lane_pts_dict)
            if tag.segment == seg.id
                push!(lanetags, tag)
            end
        end

        # sort the lanes such that the rightmost lane is lane 1
        # do this by taking the first lane,
        # then project each lane's midpoint to the perpendicular at the midpoint

        @assert(!isempty(lanetags))
        proj_positions = Array{Float64}(length(lanetags))

        first_lane_pts = lane_pts_dict[lanetags[1]]
        n = length(first_lane_pts)
        lo = first_lane_pts[div(n,2)]
        hi = first_lane_pts[div(n,2)+1]
        midpt_orig = (lo + hi)/2
        dir = polar(1.0, atan2(hi - lo) + π/2) # direction perpendicular (left) of lane

        for (i,tag) in enumerate(lanetags)
            pts = lane_pts_dict[tag]
            n = length(pts)
            midpt = (pts[div(n,2)] + pts[div(n,2)+1])/2
            proj_positions[i] = proj(midpt - midpt_orig, dir, Float64)
        end

        for (i,j) in enumerate(sortperm(proj_positions))

            tag = lanetags[j]
            boundary_left = i == length(proj_positions) ? LaneBoundary(:solid, :white) : LaneBoundary(:broken, :white)
            boundary_right = i == 1 ? LaneBoundary(:solid, :white) : LaneBoundary(:broken, :white)

            pts = lane_pts_dict[tag]
            pt_matrix = Array{Float64}(2, length(pts))
            for (k,P) in enumerate(pts)
                pt_matrix[1,k] = P.x
                pt_matrix[2,k] = P.y
            end

            println("fitting curve ", length(pts), "  "); tic()
            curve = _fit_curve(pt_matrix, desired_distance_between_curve_samples)
            toc()

            tag_new = LaneTag(seg.id, length(seg.lanes)+1)
            lane = Lane(tag_new, curve,
                        boundary_left = boundary_left,
                        boundary_right = boundary_right)
            push!(seg.lanes, lane)
            lane_new_dict[tag] = tag_new
        end
    end

    ###################################################
    # Connect the lanes
    for (tag_old, tup) in lane_next_dict
        next_pt, next_tag_old = tup
        lane = retval[lane_new_dict[tag_old]]
        next_tag_new = lane_new_dict[next_tag_old]
        dest = retval[next_tag_new]
        roadproj = proj(VecSE2(next_pt, 0.0), dest, retval)

        # println("connecting $(lane.tag) to $(dest.tag)")

        cindS = curveindex_end(lane.curve)
        cindD = roadproj.curveproj.ind

        if cindD == CURVEINDEX_START # a standard connection
            connect!(lane, dest)
            # remove any similar connection from lane_prev_dict
            if haskey(lane_prev_dict, next_tag_old) && lane_prev_dict[next_tag_old][2] == tag_old
                delete!(lane_prev_dict, next_tag_old)
            end
        else
            # otherwise connect as before
            unshift!(lane.exits,  LaneConnection(true,  cindS, RoadIndex(cindD, dest.tag)))
            push!(dest.entrances, LaneConnection(false, cindD, RoadIndex(cindS, lane.tag)))
        end
    end
    for (tag_old, tup) in lane_prev_dict
        prev_pt, prev_tag_old = tup
        lane = retval[lane_new_dict[tag_old]]
        prev_tag_new = lane_new_dict[prev_tag_old]
        prev = retval[prev_tag_new]
        roadproj = proj(VecSE2(prev_pt, 0.0), prev, retval)

        # println("connecting $(lane.tag) from $(prev.tag)")

        cindS = roadproj.curveproj.ind
        cindD = CURVEINDEX_START

        if cindS == curveindex_end(prev) # a standard connection
            @assert(!has_prev(prev))
            connect!(prev, lane)
        else
            # a standard connection
            push!(prev.exits,  LaneConnection(true,  cindS, RoadIndex(cindD, lane.tag)))
            unshift!(lane.entrances, LaneConnection(false, cindD, RoadIndex(cindS, prev.tag)))
        end
    end

    retval
end