text:
"""
Function object to map points on the cube surface to the sphere of radius r.
We need sides of length 2 centred at the origin using the approach
in http://mathproofs.blogspot.com/2005/07/mapping-cube-to-sphere.html.

Construct arguments:
  r  : sphere radius
Function object arguments:
  xyz: 3D cartesian coordinates of a point on the cube surface
"""
struct MapCubeToSphere{T} <: Function
  radius::T
end

function (map::MapCubeToSphere{T})(xyz) where T
  x,y,z = xyz
  xₛ = x*sqrt(1.0-y^2/2-z^2/2+y^2*z^2/(3.0))
  yₛ = y*sqrt(1.0-z^2/2-x^2/2+x^2*z^2/(3.0))
  zₛ = z*sqrt(1.0-x^2/2-y^2/2+x^2*y^2/(3.0))
  map.radius*Point(xₛ,yₛ,zₛ)
end

function generate_Γface_to_bgface(model)
  panel_to_entity=[23,26,24,25,21,22]
  entity_to_panel=Dict(23=>1,26=>2,24=>3,25=>4,21=>5,22=>6)
  ptr_ncells_panel=zeros(Int64,length(panel_to_entity)+1)
  labels = get_face_labeling(model)
  bgface_to_mask = Gridap.Geometry.get_face_mask(labels,"boundary",2)
  Γface_to_bgface = findall(bgface_to_mask)
  face_to_entity  = labels.d_to_dface_to_entity[3][Γface_to_bgface]
  for entity in face_to_entity
    panel=entity_to_panel[entity]
    ptr_ncells_panel[panel+1]=ptr_ncells_panel[panel+1]+1
  end
  ptr_ncells_panel[1]=1
  for i=1:length(ptr_ncells_panel)-1
    ptr_ncells_panel[i+1]=ptr_ncells_panel[i+1]+ptr_ncells_panel[i]
  end
  Γface_to_bgface_panelwise=similar(Γface_to_bgface)
  for (i,entity) in enumerate(face_to_entity)
    panel=entity_to_panel[entity]
    Γface_to_bgface_panelwise[ptr_ncells_panel[panel]]=Γface_to_bgface[i]
    ptr_ncells_panel[panel]=ptr_ncells_panel[panel]+1
  end
  Γface_to_bgface_panelwise
end

struct PolynomialMapCubedSphereDiscreteModel <: Gridap.Geometry.DiscreteModel{2,3}
  model::Gridap.Geometry.UnstructuredDiscreteModel{2,3}
end

Gridap.Geometry.get_cell_map(model::PolynomialMapCubedSphereDiscreteModel) = Gridap.Geometry.get_cell_map(model.model)
Gridap.Geometry.get_grid(model::PolynomialMapCubedSphereDiscreteModel) = Gridap.Geometry.get_grid(model.model)
Gridap.Geometry.get_grid_topology(model::PolynomialMapCubedSphereDiscreteModel) = Gridap.Geometry.get_grid_topology(model.model)
Gridap.Geometry.get_face_labeling(model::PolynomialMapCubedSphereDiscreteModel) = Gridap.Geometry.get_face_labeling(model.model)
Gridap.Geometry.Triangulation(a::PolynomialMapCubedSphereDiscreteModel) = PolynomialMapCubedSphereTriangulation(a.model)

function CubedSphereDiscreteModel(n,order; radius=1)
  function _cell_vector_to_dof_vector!(dof_vector,cell_node_ids, cell_vector)
    cache_cell_node_ids = array_cache(cell_node_ids)
    cache_cell_vector   = array_cache(cell_vector)
    for k=1:length(cell_node_ids)
       current_node_ids = getindex!(cache_cell_node_ids,cell_node_ids,k)
       current_values   = getindex!(cache_cell_vector,cell_vector,k)
       for (i,id) in enumerate(current_node_ids)
        dof_vector[current_node_ids[i]]=current_values[i]
       end
    end
  end

  domain = (-1,1,-1,1,-1,1)
  cells  = (n,n,n)
  model  = CartesianDiscreteModel(domain,cells)

  # Restrict model to cube surface
  Γface_to_bgface=generate_Γface_to_bgface(model)
  cube_surface_trian = BoundaryTriangulation(model,Γface_to_bgface)

  # Generate high-order FE map and ordering
  vector_reffe=ReferenceFE(lagrangian,VectorValue{3,Float64},order)
  V = FESpace(cube_surface_trian,vector_reffe; conformity=:H1)
  vh = interpolate(MapCubeToSphere(radius),V)
  scalar_reffe=ReferenceFE(QUAD,lagrangian,Float64,order)
  xref=Gridap.ReferenceFEs.get_node_coordinates(scalar_reffe)
  xrefₖ=Fill(xref,num_cells(cube_surface_trian))
  vhx=lazy_map(evaluate,Gridap.CellData.get_data(vh),xrefₖ)
  V = FESpace(cube_surface_trian,scalar_reffe; conformity=:H1)
  node_coordinates = Vector{Point{3,Float64}}(undef,num_free_dofs(V))
  cell_node_ids    = get_cell_dof_ids(V)
  _cell_vector_to_dof_vector!(node_coordinates,cell_node_ids,vhx)
  cell_types  = collect(Fill(1,num_cells(cube_surface_trian)))
  cell_reffes = [scalar_reffe]

  cube_surface_grid = Gridap.Geometry.UnstructuredGrid(node_coordinates,
                                                       Gridap.Arrays.Table(cell_node_ids),
                                                       cell_reffes,
                                                       cell_types,
                                                       Gridap.Geometry.Oriented())

  cube_surface_model = Gridap.Geometry.compute_active_model(cube_surface_trian)
  topology            = Gridap.Geometry.get_grid_topology(cube_surface_model)
  labeling           = Gridap.Geometry.get_face_labeling(cube_surface_model)
  PolynomialMapCubedSphereDiscreteModel(
    Gridap.Geometry.UnstructuredDiscreteModel(cube_surface_grid,topology,labeling))
end

struct AnalyticalMapCubedSphereDiscreteModel{T,B} <: Gridap.Geometry.DiscreteModel{2,3}
  cell_map::T
  cubed_sphere_linear_model::B
  function AnalyticalMapCubedSphereDiscreteModel(n;radius=1)
    domain = (-1,1,-1,1,-1,1)
    cells  = (n,n,n)
    model  = CartesianDiscreteModel(domain,cells)

    # Restrict model to cube surface
    Γface_to_bgface=generate_Γface_to_bgface(model)
    cube_surface_trian = BoundaryTriangulation(model,Γface_to_bgface)

    m1=Fill(Gridap.Fields.GenericField(MapCubeToSphere(radius)),num_cells(cube_surface_trian))
    m2=get_cell_map(cube_surface_trian)
    m=lazy_map(∘,m1,m2)

    cubed_sphere_linear_model=CubedSphereDiscreteModel(n,1)

    # Build output object
    T=typeof(m)
    B=typeof(cubed_sphere_linear_model)
    new{T,B}(m,cubed_sphere_linear_model)
  end
end

Gridap.Geometry.get_cell_map(model::AnalyticalMapCubedSphereDiscreteModel) = model.cell_map
Gridap.Geometry.get_grid(model::AnalyticalMapCubedSphereDiscreteModel) = Gridap.Geometry.get_grid(model.cubed_sphere_linear_model)
Gridap.Geometry.get_grid_topology(model::AnalyticalMapCubedSphereDiscreteModel) = Gridap.Geometry.get_grid_topology(model.cubed_sphere_linear_model)
Gridap.Geometry.get_face_labeling(model::AnalyticalMapCubedSphereDiscreteModel) = Gridap.Geometry.get_face_labeling(model.cubed_sphere_linear_model)
Gridap.Geometry.Triangulation(a::AnalyticalMapCubedSphereDiscreteModel) = AnalyticalMapCubedSphereTriangulation(a)

function CubedSphereDiscreteModel(n;radius=1)
  AnalyticalMapCubedSphereDiscreteModel(n;radius)
end

max_stars_repo_name: santiagobadia/GridapGeosciences

max_issues_repo_name: santiagobadia/GridapGeosciences

max_forks_repo_name: santiagobadia/GridapGeosciences


text:
using Random
using POMDPs
using POMDPSimulators
using DeepQLearning
using BenchmarkTools

include("../test/test_env.jl")

function evaluate(mdp, policy, rng, n_ep=100, max_steps=100)
    avg_r = 0.
    sim = RolloutSimulator(rng=rng, max_steps=max_steps)
    for i=1:n_ep
        avg_r += simulate(sim, mdp, policy)
    end
    return avg_r/=n_ep
end

function bench_prioritized_ddqn(obsdim)
    rng = MersenneTwister(1)
    mdp = TestMDP(obsdim, 4, 6)
    solver = DeepQLearningSolver(arch=QNetworkArchitecture(conv=[], fc=[32,32]),
                                 max_steps=10000, lr=0.005,
                                 eval_freq=2000,num_ep_eval=100,
                                 log_freq = 15000, verbose=false,
                                 double_q = true, dueling=true, prioritized_replay=true,
                                 rng=rng)

    policy = solve(solver, mdp)
end

function bench_drqn(obsdim)
    rng = MersenneTwister(1)
    mdp = TestMDP(obsdim, 1, 6)
    solver = DeepRecurrentQLearningSolver(arch = RecurrentQNetworkArchitecture(fc_in=[], lstm_size=32,fc_out=[32]),
                                 max_steps=10000, lr=0.005,
                                 eval_freq=2000,num_ep_eval=100,
                                 log_freq = 15000, verbose=false,
                                 double_q = true, dueling=false,
                                 rng=rng)
    policy = solve(solver, mdp)
end

for obsdim in [(5,5), (5,5,5), (20,20), (200,)]
    @show obsdim
    @info "Prioritized DDQN"
    @btime bench_prioritized_ddqn($obsdim)
    @info "DRQN"
    @btime bench_drqn($obsdim)
end

ext: jl


text:
include("setupFWItests.jl")

using Test
using LinearAlgebra

println("\tTest getData serial and parallel")
Ds,pFor  = getData(vec(m),pFor)
Dp,pForp = getData(vec(m),pForp,M2Mp,true)
for k=1:length(Dp)
	@test norm(vec(Ds[:,:,k]-fetch(Dp[k])))/norm(vec(Ds[:,:,k])) < 1e-10
end

# check sensitivities
function dFun(m,pFor,v=[])
	D,pFor = getData(m,pFor)
	if isempty(v)
		return vec(D)
	else
		dD = getSensMatVec(v,m,pFor)
		return vec(D),vec(dD)
	end
end

println("\tderivative tests")
m0 = vec(m) + .001*randn(prod(n))
f(x,v=[]) = dFun(x,pFor,v)
chkDer1, = checkDerivative(f,m0,out=false)
@test chkDer1 == true

println("\tadjoint tests")
v1 = randn(length(Ds))
v2 = randn(length(m))

t1 = dot(v1,getSensMatVec(v2,vec(m),pFor))
t2 = dot(v2,getSensTMatVec(v1,vec(m),pFor))
@test norm(t1-t2)/norm(t1)<1e-8


text:
import SparseArrays.spzeros

function inner_boundary_filter(shape)
    if shape == [3, 3]
        ibf = _inner_boundary33()
    elseif shape == [3, 3, 3]
        ibf = _inner_boundary333()
    else
        ibf = Nothing
        error("Wrong shape. Only 2D and 3D shape with size 3 is supported now.")
    end
    return ibf
end

function _inner_boundary333()
    shape = [3, 3, 3]
    prd = prod(shape)

    face_number = 0
    for i = 1:length(shape)
        # for every axis there is this number of faces
        face_number += (shape[i] + 1) * (prd // shape[i]).num
    end
    filter_arr = spzeros(Int8, face_number)

    filter_arr[14] = 1
    filter_arr[23] = 1
    filter_arr[56] = 1
    filter_arr[53] = 1
    filter_arr[90] = 1
    filter_arr[91] = 1
    return filter_arr
end

function _inner_boundary33()
    shape = [3, 3]
    filter_arr = spzeros(Int8, prod(shape))
    filter_arr[5] = 1
    filter_arr[8] = 1
    filter_arr[18] = 1
    filter_arr[19] = 1
    return filter_arr
end


text:
using HDF5, Test, DSGE, ModelConstructors

path = dirname(@__FILE__)

file = "$path/../reference/solve.h5"
TTT_expected = h5read(file, "TTT")
CCC_expected = h5read(file, "CCC")
RRR_expected = h5read(file, "RRR")

m = AnSchorfheide()
TTT, RRR, CCC = solve(m)

@testset "Check state-space system matches reference" begin
    @test @test_matrix_approx_eq TTT_expected TTT
    @test @test_matrix_approx_eq RRR_expected RRR
    @test @test_matrix_approx_eq CCC_expected CCC
end

m = Model1002("ss10")
m <= Setting(:alternative_policy, taylor93()) # check w/out using regime-switching set up
TTT_alt93, RRR_alt93, CCC_alt93 = solve(m)
m <= Setting(:regime_switching, true) # check w/using regime-switching
setup_permanent_altpol!(m, taylor93())
TTT_alt93rs, RRR_alt93rs, CCC_alt93rs = solve(m; regime_switching = true, regimes = collect(1:2), gensys_regimes = [1:2])
file = "$path/../reference/solve_alt.h5"
TTT_alt_expected93 = h5read(file, "TTT_alt93")
CCC_alt_expected93 = h5read(file, "CCC_alt93")
RRR_alt_expected93 = h5read(file, "RRR_alt93")

@testset "Check state-space system with altpolicy matches reference for Taylor93" begin
    @test @test_matrix_approx_eq TTT_alt_expected93 TTT_alt93rs[2]
    @test @test_matrix_approx_eq RRR_alt_expected93 RRR_alt93rs[2]
    @test @test_matrix_approx_eq CCC_alt_expected93 CCC_alt93rs[2]
    @test @test_matrix_approx_eq TTT_alt_expected93 TTT_alt93
    @test @test_matrix_approx_eq RRR_alt_expected93 RRR_alt93
    @test @test_matrix_approx_eq CCC_alt_expected93 CCC_alt93
    @test maximum(abs.(TTT_alt_expected93 - TTT_alt93rs[1])) > 1e-6
    @test maximum(abs.(RRR_alt_expected93 - RRR_alt93rs[1])) > 1e-6
end

m = Model1002("ss10")
m <= Setting(:alternative_policy, taylor99())
TTT_alt99, RRR_alt99, CCC_alt99 = solve(m)
m <= Setting(:regime_switching, true)
setup_permanent_altpol!(m, taylor99())
TTT_alt99rs, RRR_alt99rs, CCC_alt99rs = solve(m; regime_switching = true, regimes = collect(1:2), gensys_regimes = [1:2])
file = "$path/../reference/solve_alt.h5"
TTT_alt_expected99 = h5read(file, "TTT_alt99")
CCC_alt_expected99 = h5read(file, "CCC_alt99")
RRR_alt_expected99 = h5read(file, "RRR_alt99")

@testset "Check state-space system with altpolicy matches reference for Taylor99" begin
    @test @test_matrix_approx_eq TTT_alt_expected99 TTT_alt99rs[2]
    @test @test_matrix_approx_eq RRR_alt_expected99 RRR_alt99rs[2]
    @test @test_matrix_approx_eq CCC_alt_expected99 CCC_alt99rs[2]
    @test @test_matrix_approx_eq TTT_alt_expected99 TTT_alt99
    @test @test_matrix_approx_eq RRR_alt_expected99 RRR_alt99
    @test @test_matrix_approx_eq CCC_alt_expected99 CCC_alt99
    @test maximum(abs.(TTT_alt_expected99 - TTT_alt99rs[1])) > 1e-6
    @test maximum(abs.(RRR_alt_expected99 - RRR_alt99rs[1])) > 1e-6
end

nothing


text:
# TXTFile.jl

import Base: ==, getindex
import Base: length, filesize, start, next, done, parse, open, push!, keys, values, map, endof
export TXTFile, save

type TXTFile <: AbstractFile
  path::AbstractString
  data::Union{AbstractArray,Void}
  TXTFile(path) = new(path, nothing)
end

type TXTFileIO
  file::TXTFile
  isdirty::Bool
end

function support(filename::AbstractString, ::Type{Val{symbol(".txt")}})
  TXTFile(filename)
end

function ==(a::TXTFile, b::TXTFile)
  a.path == b.path
end

function parse(file::TXTFile)
  if nothing == file.data
    f = open(file.path)
    file.data = map(rstrip, readlines(f))
    close(f)
  end
  file.data
end

function open(fn::Function, file::TXTFile)
  parse(file)
  io = TXTFileIO(file, false)
  fn(io)
  if io.isdirty
    save(file, join(io.file.data, '\n'))
  end
end

function save(file::TXTFile, data)
  f = open(file.path, "w")
  write(f, data)
  close(f)
end

function push!(file::TXTFile, data::Any)
  open(file) do io
    push!(io, data)
  end
end

function push!(io::TXTFileIO, data::Any)
  push!(io.file.data, data)
  io.isdirty = true
end

map(fn::Function, file::TXTFile) = map(fn, parse(file))

length(file::TXTFile) = length(parse(file))
filesize(file::TXTFile) = filesize(file.path)

# indexing
getindex(file::TXTFile, a::Any) = getindex(parse(file), a)

# iterators
start(f::TXTFile) = start(parse(f))
next(file::TXTFile, i) = next(parse(file), i)
done(file::TXTFile, i) = done(parse(file), i)
endof(file::TXTFile) = endof(parse(file))

# dict
keys(file::TXTFile) = collect(keys(parse(file)))
values(file::TXTFile) = collect(values(parse(file)))


text:
@enum(JoystickAxis,
    joystickX,    #< The X axis
    joystickY,    #< The Y axis
    joystickZ,    #< The Z axis
    joystickR,    #< The R axis
    joystickU,    #< The U axis
    joystickV,    #< The V axis
    joystickPovX, #< The X axis of the point-of-view hat
    joystickPovY  #< The Y axis of the point-of-view hat
)

function is_connected(joystick::Integer)
    ccall((:sfJoystick_isConnected, libcsfml_window), Bool, (UInt32,), joystick)
end

function get_button_count(joystick::Integer)
    ccall((:sfJoystick_getButtonCount, libcsfml_window), UInt32, (UInt32,), joystick)
end

function has_axis(joystick::Integer, axis::JoystickAxis)
    ccall((:sfJoystick_hasAxis, libcsfml_window), Bool, (UInt32, Int32,), joystick, axis)
end

function is_button_pressed(joystick::Integer, button::Integer)
    ccall((:sfJoystick_isButtonPressed, libcsfml_window), Bool, (UInt32, UInt32,), joystick, button)
end

function get_axis_position(joystick::Integer, axis::JoystickAxis)
    ccall((:sfJoystick_getAxisPosition, libcsfml_window), Float32, (UInt32, Int32,), joystick, axis)
end

function joystick_update()
    ccall((:sfJoystick_update, libcsfml_window), Void, ())
end


text:
module Nanoid

using Random

const urlalphabet = "useandom-26T198340PX75pxJACKVERYMINDBUSHWOLF_GQZbfghjklqvwyzrict"

rng = Random.RandomDevice()

function randomcharacter(from::String)::Char
    return rand(rng, from)
end

function generate(alphabet::String, length::Integer)::String
    id = []

    for _ = 1:length
        randomchar = string(randomcharacter(alphabet))
        append!(id, randomchar)
    end

    join(id, "")
end

"""
    nanoid(size::Integer = 21; alphabet::String = "") -> String

Generate a random Nanoid with the given size and alphabet.
Size defaults to 21, and alphabet defaults to the URL-safe alphabet.

## Examples

```jldoctest
julia> nanoid()
"tcDIFxMJHw6UgUnbI_6za"

julia> nanoid(10)
"OutzQWa1SB"

julia> nanoid(; alphabet = "ABCDEFGHIJ")
"GEHCAGFJIAHCJGIHHIHFF"

julia> nanoid(10; alphabet = "abcdef")
"abeafebaab"
```
"""
function nanoid(size::Integer=21; alphabet::String = "")::String
    if size < 1
        error("size must be at least 1")
    end

    if alphabet == ""
        alphabet = urlalphabet
    end

    generate(alphabet, size)
end

export nanoid

end


text:
if isfile(joinpath(dirname(dirname(@__FILE__)), "..", "deps", "deps.jl"))
    include("../../deps/deps.jl")
else
    ext = Sys.iswindows() ? ".dll" : ".so"
    println(joinpath(dirname(dirname(@__FILE__)), "deps", "deps.jl"))
    error("lib_nes_env" * ext *  " not properly installed. Please run Pkg.build(\"ArcadeLearningEnvironment\")")
end

include("ROMInterface.jl")

NESEnvPtr = Ptr{Cvoid}
ControllerVectorPtr = Ptr{UInt8}
ScreenTensorPtr = Ptr{UInt8}
RAMVectorPtr = Ptr{UInt8}

#--------------------------- lib_nes_env Interface -----------------------------

# Returns screen width and height.
Width()::Int  = ccall((:Width, lib_nes_env), Cuint, ())
Height()::Int = ccall((:Height, lib_nes_env), Cuint, ())

# Returns the address of the controller for `nesenvptr` at port number `port`.
Controller(nesenvptr::NESEnvPtr, port::Integer) = ccall((:Controller, lib_nes_env),
                                                ControllerVectorPtr, (NESEnvPtr, Cuint), nesenvptr, Cuint(port))

# Returns the pointer to the 1D array that contains the RGB screen picture.
Screen(nesenvptr::NESEnvPtr) = ccall((:Screen, lib_nes_env), ScreenTensorPtr, (NESEnvPtr,), nesenvptr)

# Returns the address of the RAM
Memory(nesenvptr::NESEnvPtr) = ccall((:Memory, lib_nes_env), RAMVectorPtr, (NESEnvPtr,), nesenvptr)

# Resets the game
Reset(nesenvptr::NESEnvPtr) = ccall((:Reset, lib_nes_env), Cvoid, (NESEnvPtr,), nesenvptr)

# Performs a step on the emulator
Step(nesenvptr::NESEnvPtr) = ccall((:Step, lib_nes_env), Cvoid, (NESEnvPtr,), nesenvptr)

# Backup of the emulator state
Backup(nesenvptr::NESEnvPtr) = ccall((:Backup, lib_nes_env), Cvoid, (NESEnvPtr,), nesenvptr)

# Restore the backed up state into the NES emulator
Restore(nesenvptr::NESEnvPtr) = ccall((:Restore, lib_nes_env), Cvoid, (NESEnvPtr,), nesenvptr)

# Close the environment
Close(nesenvptr::NESEnvPtr) = ccall((:Close, lib_nes_env), Cvoid, (NESEnvPtr,), nesenvptr)

#---------------------------------Utilites--------------------------------------

const h = Height()
const w = Width()
const SCREEN_SHAPE_32_BIT = (h, w, 4)
const SCREEN_SHAPE_24_BIT = (h, w, 3)

"""
    Initialize(rom_path::String)

Creates a new NES emulator instance and loads the game at `rom_path`
"""
function Initialize(rom_path::AbstractString)
    edited_rom = push!(Cwchar_t[char for char in rom_path], Cwchar_t(0))
    ccall((:Initialize, lib_nes_env), NESEnvPtr, (Ptr{Cwchar_t},), edited_rom)
end

"""
    screen_buffer(nesenvptr::NESEnvPtr)

Setup the screen buffer from the C++ code
"""
function screen_buffer(nesenvptr::NESEnvPtr)
    address = Screen(nesenvptr)
    buffer = unsafe_wrap(Array, address, prod(SCREEN_SHAPE_32_BIT))
    # create a screen tensor by reshaping the contents of the buffer and flipping
    # the final set of bytes for little endian machines, which it likely is

    # Just in case, if the machine is big endian, reverse what we did...
    #Base.ENDIAN_BOM == 0x01020304 && (screen_tensor = screen_tensor[:, :, end:-1:1])

    return buffer
end

"""
    ram_vector(nesenvptr::NESEnvPtr)

Setup the RAM buffer from the C++ code
"""
function ram_buffer(nesenvptr::NESEnvPtr)
    # get the address of the RAM
    address = Memory(nesenvptr)
    # create a buffer from the contents of the address location
    buffer = unsafe_wrap(Array, address, 0x800)
    return buffer
end

"""
    controller_buffer(nesenvptr::NESEnvPtr, port)

Find the pointer to a controller and setup a buffer. Accepts `port` which is the
port of the controller and returns the buffer with the controller's binary data.
"""
function controller_buffer(nesenvptr::NESEnvPtr, port)
    # Get address of the controller
    address = Controller(nesenvptr, port)
    # Create a memory buffer using the address pointer.
    buffer = unsafe_wrap(Array, address, 1)
    return buffer
end

#-----------------------NESEnv functions----------------------------------------

mutable struct NESEnv
    envptr::Union{NESEnvPtr, Nothing}
    rom_path::String
    has_backup::Union{Nothing, Bool}
    done::Bool
    controllers::Array{Array{UInt8, N}, 1} where N
    screen_buf::Array{UInt8, 1}
    ram::Array{UInt8, 1}
end

function NESEnv(rom_path::AbstractString)
    if isfile(rom_path)
        rom_path = rom_path
    elseif isfile(joinpath(@__DIR__, "..", "..", "deps", "roms", rom_path * ".nes"))
        rom_path = joinpath(@__DIR__, "..", "..", "deps", "roms", rom_path * ".nes")
    else
        error("File at path $rom_path not found.")
    end

    rom = ROM(rom_path)
    # Ensure that PRG-ROM is present and no trainers are present
    rom.prg_rom_size == 0 &&
        error("ROM has no PRG-ROM banks. No can do.")
    rom.has_trainer &&
        error("ROM has trainer. Trainer isn't supported.")

    # Try to read the PRG ROM and raises an error if it fails
    _ = rom.prg_rom
    # Try to read the CHR ROM and raise an error if it fails
    _ = rom.chr_rom

    # check the TV system and/or if mapper is implemented
    if rom.is_pal
        error("ROM is PAL. PAL is not supported")
    elseif rom.mapper ∉ (0, 1, 2, 3)
        error("ROM has unsupported mapper number $(mapper(rom))")
    end

    # Initialize the C++ object for running the environment
    envptr = Initialize(rom_path)
    Reset(envptr)
    has_backup = false
    done = true
    controllers = [controller_buffer(envptr, port) for port in 0:1]
    screen_buf = screen_buffer(envptr)
    ram = ram_buffer(envptr)

    NESEnv(envptr, rom_path, has_backup, done, controllers, screen_buf, ram)
end

function get_obs(nesenv::NESEnv)
    r_array = reshape(nesenv.screen_buf[1:4:end], w, h)
    g_array = reshape(nesenv.screen_buf[2:4:end], w, h)
    b_array = reshape(nesenv.screen_buf[3:4:end], w, h)

    rgb_array = cat(r_array, g_array, b_array; dims=3)
end

function Base.getproperty(obj::NESEnv, sym::Symbol)
    if sym == :screen
        return get_obs(obj)
    else
        return Base.getfield(obj, sym)
    end
end

"""
    frame_advance!(nesenv::NESEnv, action)

Advance a fram in the emulator with an action.
"""
function frame_advance!(nesenv::NESEnv, action)
    nesenv.controllers[1][:] .= action
    Step(nesenv.envptr)
end

"""
    backup(nesenv::NESEnv)

Backup the NES state in the emulator
"""
function backup!(nesenv::NESEnv)
    Backup(nesenv.envptr)
    nesenv.has_backup = true
end

"""
    restore(nesenv::NESEnv)

Restore the backup state into the NES emulator.
"""
restore!(nesenv::NESEnv) = Restore(nesenv.envptr)

"""
    reset(nesenv::NESEnv)

Reset the state of the environment and returns an initial observation
"""
function reset!(nesenv::NESEnv)
    nesenv.has_backup ? restore(nesenv) : Reset(nesenv.envptr)
    nesenv.done = false
end

"""
    step!(nesenv::NESEnv, action)

Run one frame of the NES and return the relevant observation data.
"""
function step!(nesenv::NESEnv, action)
    # Raise an error if the environment is done.
    nesenv.done &&
        error("Cannot step in a done environment! Call `reset!`")

    # Set the action on the controller.
    self.controllers[1][:] = action
    # pass the action to the emulator as an unsigned byte
    Step(nesenv.envptr)

    reward = get_reward(nesenv)
    nesenv.done = get_done(nesenv)
    info = get_info(nesenv)
    did_step(nesenv)

    reward = clamp(reward, nesenv.reward_range[1], nesenv.reward_range[2])

    return nesenv.screen, reward, nesenv.done, info
end

"""
    close(nesenv::NESEnv)

Close the environment.
"""
function close(nesenv::NESEnv)
    isnothing(nesenv.envptr) && error("env has already been closed.")
    Close(nesenv.envptr)
    nesenv.envptr = nothing
end

"""
    get_keys_to_action(nesenv::NESEnv)

Return the dictionary of keyboard keys to actions.
"""
function get_keys_to_action(nesenv::NESEnv)
    buttons = Int[
        'd', # right
        'a', # left
        's', # down
        'w', # up
       '\r', # start
        ' ', # stop
        'p', # B
        'o'  # A
    ]
    keys_to_action = Dict()

    for i=0:255
        byte = string(i, base=2, pad=8)
        pressed = sort(buttons[Bool[parse(Int, num) for num in byte]])
        keys_to_action[pressed] = byte
    end
    return keys_to_action
end

#println(NESEnv("/Users/kartikeygupta/Workspace/Julia/Gym/NES/deps/roms/smb.nes"))


text:
"""
    WebSockets
This module implements the WebSockets protocol. It relies on the package HTTP.jl.

Websocket|server relies on a client initiating the connection.
Websocket|client initiate the connection.

The client side of the connection is most typically a browser with
scripts enabled. Browsers are always the initiating, client side. But the
peer can be any program, in any language, that follows the protocol. That
includes another Julia session, running in a parallel process or task.

    Future improvements:
1. Check rsv1 to rsv3 values. This will reduce bandwidth.
2. Optimize maskswitch!, possibly threaded above a certain limit.
3. Split messages over several frames.
"""
module WebSockets
using Dates
using Logging
import Sockets
import Sockets: TCPSocket,        # For locked_write, show
                IPAddr,           # For serve
                InetAddr          # For serve
import Base64:  base64decode,     # For generate_websocket_key
                base64encode      # For open client websocket
import Base:    IOServer,         # For serve
                ReinterpretArray, # For data type
                buffer_writes,    # For init_socket
                CodeUnits,        # For data type
                throwto           # For readframe_nonblocking
import HTTP                       # Depend on WebSockets.HTTP only
                                  # to avoid version conflicts!
import HTTP.Servers.MbedTLS       # For further imports
import HTTP.Servers.MbedTLS:
                MD_SHA1,          # For generate_websocket_key
                digest            # For generate_websocket_key

# further imports from HTTP in this file
include("HTTP.jl")

# A logger based on ConsoleLogger. This has an effect
# only if the user chooses to use WebSocketLogger.
include("Logger/websocketlogger.jl")

export WebSocket,
       serve,
       readguarded,
       writeguarded,
       write,
       read,
       close,
       subprotocol,
       target,
       send_ping,
       send_pong,
       WebSocketClosedError,
       addsubproto,
       WebSocketLogger,
       @wslog,
       Wslog

# revisit the need for defining this union type for method definitions. The functions would
# probably work just as fine with duck typing.
const Dt = Union{ReinterpretArray{UInt8,1,UInt16,Array{UInt16,1}},
            Vector{UInt8},
            CodeUnits{UInt8,String}   }
"A reasonable amount of time"
const TIMEOUT_CLOSEHANDSHAKE = 10.0

@enum ReadyState CONNECTED=0x1 CLOSING=0x2 CLOSED=0x3

""" Buffer writes to socket till flush (sock)"""
init_socket(sock) = buffer_writes(sock)

struct WebSocketClosedError <: Exception
    message::String
end

struct WebSocketError <: Exception
    status::Int16
    message::String
end

"Status codes according to RFC 6455 7.4.1"
const codeDesc = Dict{Int, String}(
    1000=>"Normal",                     1001=>"Going Away",
    1002=>"Protocol Error",             1003=>"Unsupported Data",
    1004=>"Reserved",                   1005=>"No Status Recvd- reserved",
    1006=>"Abnormal Closure- reserved", 1007=>"Invalid frame payload data",
    1008=>"Policy Violation",           1009=>"Message too big",
    1010=>"Missing Extension",          1011=>"Internal Error",
    1012=>"Service Restart",            1013=>"Try Again Later",
    1014=>"Bad Gateway",                1015=>"TLS Handshake")

"""
A WebSocket is a wrapper over a TCPSocket. It takes care of wrapping outgoing
data in a frame and unwrapping (and concatenating) incoming data.
"""
mutable struct WebSocket{T <: IO} <: IO
    socket::T
    server::Bool
    state::ReadyState

    function WebSocket{T}(socket::T,server::Bool) where T
        init_socket(socket)
        new(socket, server, CONNECTED)
    end
end
WebSocket(socket,server) = WebSocket{typeof(socket)}(socket,server)

# WebSocket Frames
#
#      0                   1                   2                   3
#      0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
#     +-+-+-+-+-------+-+-------------+-------------------------------+
#     |F|R|R|R| opcode|M| Payload len |    Extended payload length    |
#     |I|S|S|S|  (4)  |A|     (7)     |             (16/64)           |
#     |N|V|V|V|       |S|             |   (if payload len==126/127)   |
#     | |1|2|3|       |K|             |                               |
#     +-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
#     |     Extended payload length continued, if payload len == 127  |
#     + - - - - - - - - - - - - - - - +-------------------------------+
#     |                               |Masking-key, if MASK set to 1  |
#     +-------------------------------+-------------------------------+
#     | Masking-key (continued)       |          Payload Data         |
#     +-------------------------------- - - - - - - - - - - - - - - - +
#     :                     Payload Data continued ...                :
#     + - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
#     |                     Payload Data continued ...                |
#     +---------------------------------------------------------------+
#

# Opcode values
""" *  %x0 denotes a continuation frame"""
const OPCODE_CONTINUATION = 0x00
""" *  %x1 denotes a text frame"""
const OPCODE_TEXT = 0x1
""" *  %x2 denotes a binary frame"""
const OPCODE_BINARY = 0x2
#  *  %x3-7 are reserved for further non-control frames
#
""" *  %x8 denotes a connection close"""
const OPCODE_CLOSE = 0x8
""" *  %x9 denotes a ping"""
const OPCODE_PING = 0x9
""" *  %xA denotes a pong"""
const OPCODE_PONG = 0xA
# *  %xB-F are reserved for further control frames

"""
Handshakes with subprotocols are rejected by default.
Add to acceptable SUBProtocols through e.g.
```julia
   addsubproto("json")
```
Also see function subprotocol
"""
const SUBProtocols= Array{String,1}()

"""
    write_fragment(io, islast, opcode, hasmask, data::Array{UInt8})
Write the raw frame to a bufffer. Websocket|client must set 'hasmask'.
"""
function write_fragment(io::IO, islast::Bool, opcode, hasmask::Bool, data::Dt)
    l = length(data)
    b1::UInt8 = (islast ? 0b1000_0000 : 0b0000_0000) | opcode

    mask::UInt8 = hasmask ? 0b1000_0000 : 0b0000_0000

    write(io, b1)
    if l <= 125
        write(io, mask | UInt8(l))
    elseif l <= typemax(UInt16)
        write(io, mask | UInt8(126))
        write(io, hton(UInt16(l)))
    elseif l <= typemax(UInt64)
        write(io, mask | UInt8(127))
        write(io, hton(UInt64(l)))
    else
        error("Attempted to send too much data for one websocket fragment\n")
    end
    if hasmask
        if opcode == OPCODE_TEXT
            # Avoid masking Strings bytes in place.
            # This makes client websockets slower than server websockets.
            data = copy(data)
        end
        # Write the random masking key to io, also mask the data in-place
        write(io, maskswitch!(data))
    end
    write(io, data)
end

""" Write without interruptions"""
function locked_write(io::IO, islast::Bool, opcode, hasmask::Bool, data::Dt)
    isa(io, TCPSocket) && lock(io.lock)
    try
        write_fragment(io, islast, opcode, hasmask, data)
    finally
        if isa(io, TCPSocket)
            flush(io)
            unlock(io.lock)
        end
    end
end

""" Write text data; will be sent as one frame."""
function Base.write(ws::WebSocket,data::String)
    # add a method for reinterpreted strings as well? See const Dt.
  #  locked_write(ws.socket, true, OPCODE_TEXT, !ws.server, Vector{UInt8}(data)) # Vector{UInt8}(String) will give a warning in v0.7.
  locked_write(ws.socket, true, OPCODE_TEXT, !ws.server, codeunits(data)) # Vector{UInt8}(String) will give a warning in v0.7.
end

""" Write binary data; will be sent as one frame."""
function Base.write(ws::WebSocket, data::Array{UInt8})
    locked_write(ws.socket, true, OPCODE_BINARY, !ws.server, data)
end

function write_ping(io::IO, hasmask, data = UInt8[])
    locked_write(io, true, OPCODE_PING, hasmask, data)
end
""" Send a ping message, optionally with data."""
send_ping(ws, data...) = write_ping(ws.socket, !ws.server, data...)

function write_pong(io::IO, hasmask, data = UInt8[])
    locked_write(io, true, OPCODE_PONG, hasmask, data)
end
""" Send a pong message, optionally with data."""
send_pong(ws, data...) = write_pong(ws.socket, !ws.server, data...)

"""
    close(ws::WebSocket)
    close(ws::WebSocket, statusnumber = n)
    close(ws::WebSocket, statusnumber = n, freereason = "my reason")
Send an OPCODE_CLOSE frame, and wait for the same response or until
a reasonable amount of time, $(round(TIMEOUT_CLOSEHANDSHAKE, digits=1)) s, has passed.
Data received while closing is dropped.
Status number n according to RFC 6455 7.4.1 can be included, see WebSockets.codeDesc
"""
function Base.close(ws::WebSocket; statusnumber = 0, freereason = "")
    if isopen(ws)
        ws.state = CLOSING
        if statusnumber == 0
            locked_write(ws.socket, true, OPCODE_CLOSE, !ws.server, UInt8[])
        elseif freereason == ""
            statuscode = reinterpret(UInt8, [hton(UInt16(statusnumber))])
            locked_write(ws.socket, true, OPCODE_CLOSE, !ws.server, copy(statuscode))
        else
            statuscode = vcat(reinterpret(UInt8, [hton(UInt16(statusnumber))]),
                                Vector{UInt8}(freereason))
            locked_write(ws.socket, true, OPCODE_CLOSE, !ws.server, copy(statuscode))
        end

        # Wait till the peer responds with an OPCODE_CLOSE while discarding any
        # trailing bytes received.
        #
        # We have no guarantee that the peer is actually reading our OPCODE_CLOSE
        # frame. If not, the peer's state will not change, and we will not receive
        # an aknowledgment of closing. We use a nonblocking read and give up
        # after TIMEOUT_CLOSEHANDSHAKE
        #
        # This process is
        # complicated by potential blocking reads on the WebSocket in other Tasks
        # which may receive the response control frame. Synchronization of who is
        # responsible for closing the underlying socket is done using the
        # WebSocket's state. When this side initiates closing the connection it is
        # responsible for cleaning up, when the other side initiates the close the
        # read method is.
        #
        # The exception handling is necessary as read_frame will error when the
        # OPCODE_CLOSE control frame is received by a potentially blocking read in
        # another Task
        #
        try
            t1 = time() + TIMEOUT_CLOSEHANDSHAKE
            while isopen(ws) && time() < t1
                wsf = readframe_nonblocking(ws)
                if is_control_frame(wsf) && (wsf.opcode == OPCODE_CLOSE)
                    ws.state = CLOSED
                end
            end
            if isopen(ws.socket)
                close(ws.socket)
            end
        catch err
            # Typical 'errors' received while closing down are neglected.
            # Unknown errors are rethrown.
            errtyp = typeof(err)
            errtyp != InterruptException &&
                errtyp != Base.IOError &&
                errtyp != HTTP.IOExtras.IOError &&
                errtyp != Base.BoundsError &&
                errtyp != Base.EOFError &&
                errtyp != Base.ArgumentError &&
                rethrow(err)
        end
    else
        ws.state = CLOSED
    end
end

"""
    isopen(::WebSocket)-> Bool
A WebSocket is closed if the underlying TCP socket closes, or if we send or
receive a close message.
"""
Base.isopen(ws::WebSocket) = (ws.state != CLOSED) && isopen(ws.socket)

Base.eof(ws::WebSocket) = (ws.state == CLOSED) || eof(ws.socket)

""" Represents one (received) message frame."""
mutable struct WebSocketFragment
    is_last::Bool
    rsv1::Bool     # Set for compressed messages.
    rsv2::Bool
    rsv3::Bool
    opcode::UInt8  # This is actually a UInt4 value.
    is_masked::Bool
    payload_len::UInt64
    maskkey::Vector{UInt8}  # This will be 4 bytes on frames from the client.
    data::Vector{UInt8}  # For text messages, this is a String.
end

""" This constructor handles conversions from bytes to bools."""
function WebSocketFragment(
     fin::UInt8
    , rsv1::UInt8
    , rsv2::UInt8
    , rsv3::UInt8
    , opcode::UInt8
    , masked::UInt8
    , payload_len::UInt64
    , maskkey::Vector{UInt8}
    , data::Vector{UInt8})

    WebSocketFragment(
      fin != 0
    , rsv1 != 0
    , rsv2 != 0
    , rsv3 != 0
    , opcode
    , masked != 0
    , payload_len
    , maskkey
    , data)
end

""" Control frames have opcodes with the highest bit = 1."""
is_control_frame(msg::WebSocketFragment) = (msg.opcode & 0b0000_1000) > 0

""" Respond to pings, ignore pongs, respond to close."""
function handle_control_frame(ws::WebSocket, wsf::WebSocketFragment)
    if wsf.opcode == OPCODE_CLOSE
        ws.state = CLOSED
        try
            locked_write(ws.socket, true, OPCODE_CLOSE, !ws.server, UInt8[])
        catch e
        end
        # Find out why the other side wanted to close.
        # RFC 6455 5.5.1. If there is a status code, it's a two-byte number in network order.
        if wsf.payload_len == 0
            reason = " No reason "
        elseif wsf.payload_len == 2
            scode = Int(reinterpret(UInt16, reverse(wsf.data))[1])
            reason = string(scode) * ": " * get(codeDesc, scode, "")
        else
            scode = Int(reinterpret(UInt16, reverse(wsf.data[1:2]))[1])
            reason = string(scode) * ": " * String(wsf.data[3:end])
        end
        throw(WebSocketClosedError("ws|$(ws.server ? "server" : "client") respond to OPCODE_CLOSE " * reason))
    elseif wsf.opcode == OPCODE_PING
        @debug ws, " received OPCODE_PING"
        send_pong(ws, wsf.data)
    elseif wsf.opcode == OPCODE_PONG
        @debug ws, " received OPCODE_PING"
        # Nothing to do here; no reply is needed for a pong message.
    else  # %xB-F are reserved for further control frames
        error("while handle_control_frame(ws|$(ws.server ? "server" : "client"), wsf): Unknown opcode $(wsf.opcode)")
    end
end

""" Read a frame: turn bytes from the websocket into a WebSocketFragment."""
function read_frame(ws::WebSocket)
    # Try to read two bytes. There is no guarantee that two bytes are actually allocated.
    ab = Array{UInt8}(undef, 2)
    if readbytes!(ws.socket, ab) != 2
      throw(WebSocketError(1006, "Client side closed socket connection"))
    end

    #=
    Browsers will seldom close in the middle of writing to a socket,
    but other clients often do, and the stacktraces can be very long.
    ab can be assigned, but of length 1. Use an enclosing try..catch in the calling function
    =#
    a = ab[1]
    fin    = (a & 0b1000_0000) >>> 7  # If fin, then is final fragment
    rsv1   = a & 0b0100_0000  # If not 0, fail.
    rsv2   = a & 0b0010_0000  # If not 0, fail.
    rsv3   = a & 0b0001_0000  # If not 0, fail.
    opcode = a & 0b0000_1111  # If not known code, fail.

    b = ab[2]
    mask = (b & 0b1000_0000) >>> 7
    hasmask = mask != 0

    if hasmask != ws.server
        if ws.server
            msg = "WebSocket|server cannot handle incoming messages without mask. Ref. rcf6455 5.3"
        else
            msg = "WebSocket|client cannot handle incoming messages with mask. Ref. rcf6455 5.3"
        end
        throw(WebSocketError(1002, msg))
    end

    payload_len::UInt64 = b & 0b0111_1111
    if payload_len == 126
        payload_len = ntoh(read(ws.socket, UInt16))  # 2 bytes
    elseif payload_len == 127
        payload_len = ntoh(read(ws.socket, UInt64))  # 8 bytes
    end

    maskkey = hasmask ? read(ws.socket, 4) : UInt8[]

    data = read(ws.socket,Int(payload_len))
    hasmask && maskswitch!(data, maskkey)

    return WebSocketFragment(
        fin, rsv1, rsv2, rsv3,
        opcode, mask, payload_len, maskkey, data
    )
end

"""
    read(ws::WebSocket)
Typical use:
    msg = String(read(ws))
Read one non-control message from a WebSocket. Any control messages that are
read will be handled by the handle_control_frame function.
Only the data (contents/body/payload) of the message will be returned as a
Vector{UInt8}.

This function will not return until a full non-control message has been read.
"""
function Base.read(ws::WebSocket)
    if !isopen(ws)
        error("Attempt to read from closed WebSocket|$(ws.server ? "server" : "client"). First isopen(ws), or use readguarded(ws)!")
    end
    try
        frame = read_frame(ws)
        # Handle control (non-data) messages.
        if is_control_frame(frame)
            # Don't return control frames; they're not interesting to users.
            handle_control_frame(ws, frame)
            # Recurse to return the next data frame.
           return read(ws)
        end

        # Handle data message that uses multiple fragments.
        if !frame.is_last
            return vcat(frame.data, read(ws))
        end
        return frame.data
    catch err
        try
            server_str = ws.server ? "server" : "client"
            if err isa InterruptException
                msg = "while read(ws|$(server_str) received InterruptException."
                # This exception originates from this side. Follow close protocol so as not to irritate the other side.
                close(ws, statusnumber = 1006, freereason = msg)
                throw(WebSocketClosedError(msg * " Performed closing handshake."))
            elseif err isa WebSocketError
                # This exception originates on the other side. Follow close protocol with reason.
                close(ws, statusnumber = err.status)
                throw(WebSocketClosedError("while read(ws|$(server_str)) $(err.message) - Performed closing handshake."))
            elseif err isa Base.IOError || err isa Base.EOFError
                throw(WebSocketClosedError("while read(ws|$(server_str)) $(string(err))"))
            else
                # Unknown cause, give up continued execution.
                # If this happens in a multiple fragment message, the accumulated
                # stacktrace could be very long since read(ws) is iterative.
                rethrow(err)
            end
        finally
            if isopen(ws.socket)
                close(ws.socket)
            end
            ws.state = CLOSED
        end
    end
    return UInt8[]
end

"""
For the closing handshake, we won't wait indefinitely for non-responsive clients.
Returns a throwaway frame if the socket happens to be empty
"""
function readframe_nonblocking(ws)
    chnl= Channel{WebSocketFragment}(1)
    # Read, output put to Channel for type stability
    function _readinterruptable(c::Channel{WebSocketFragment})
        try
            put!(chnl, read_frame(ws))
        catch
            # Output a dummy frame that is not a control frame.
            put!(chnl, WebSocketFragment(false, false, false, false,
                                UInt8(0), false, UInt64(0),
                                Vector{UInt8}([0x0,0x0,0x0,0x0]),
                                Vector{UInt8}()))
        end
    end
    # Start reading as a task. Will not return if there is nothing to read
    rt = @async _readinterruptable(chnl)
    bind(chnl, rt)
    yield()
    # Define a task for throwing interrupt exception to the (possibly blocked) read task.
    # We don't start this task because it would never return
    killta = @task try
        throwto(rt, InterruptException())
    catch
    end
    # We start the killing task. When it is scheduled the second time,
    # we pass an InterruptException through the scheduler.
    try
        schedule(killta, InterruptException(), error = false)
    catch
    end
    # We now have content on chnl, and no additional tasks.
    take!(chnl)
end

"""
    WebSocket Handshake Procedure
`generate_websocket_key(key)` transforms a websocket client key into the server's accept
value. This is done in three steps:
1. Concatenate key with magic string from RFC.
2. SHA1 hash the resulting base64 string.
3. Encode the resulting number in base64.
This function then returns the string of the base64-encoded value.
"""
function generate_websocket_key(key)
    hashkey = "$(key)258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
    return base64encode(digest(MD_SHA1, hashkey))
end

"""
    maskswitch!(data)
    maskswitch!(data, key:: 4-element Vector{UInt8})

Masks or unmasks data in-place, returns the key used.
Calling twice with the same key restores data.
Ref. RFC 6455 5-3.
"""
function maskswitch!(data, mask = rand(UInt8, 4))
    for i in 1:length(data)
        data[i] = data[i] ⊻ mask[((i-1) % 4)+1]
    end
    return mask
end

"Used in handshake. See SUBProtocols"
hasprotocol(s::AbstractString) = in(s, SUBProtocols)

"Used to specify acceptable subprotocols. See SUBProtocols"
function addsubproto(name)
    push!(SUBProtocols, string(name))
    return true
end

"""
`target(request) => String`

Convenience function for reading upgrade request target.
    E.g.
```julia
    function gatekeeper(req, ws)
        if target(req) == "/gamepad"
            @spawnat 2 gamepad(ws)
        elseif target(req) == "/console"
            @spawnat 3 console(ws)
            ...
        end
    end
```
Then, in browser javascript (or equivalent with Julia WebSockets.open( , ))
```javascript
function load(){
    var wsuri = document.URL.replace("http:", "ws:");
    ws1 = new WebSocket(wsuri + "/gamepad");
    ws2 = new WebSocket(wsuri + "/console");
    ws3 = new WebSocket(wsuri + "/graphics");
    ws4 = new WebSocket(wsuri + "/audiochat");
    ws1.onmessage = function(e){vibrate(e.data)}
    } // load

```
"""
function target   # Methods added in include files
end

"""
`subprotocol(request) => String`

Convenience function for reading upgrade request subprotocol.
Acceptable subprotocols need to be predefined using
addsubproto(myprotocol). No other subprotocols will pass the handshake.
E.g.
```julia
WebSockets.addsubproto("instructions")
WebSockets.addsubproto("relay_backend")
function gatekeeper(req, ws)
    subpr = WebSockets.subprotocol(req)
    if subpr == "instructions"
        instructions(ws)
    elseif subpr == "relay_backend"
        relay_backend(ws)
    end
end
```

Then, in browser javascript (or equivalent with Julia WebSockets.open( , ))
```javascript
function load(){
    var wsuri = document.URL.replace("http:", "ws:");
    ws1 = new WebSocket(wsuri, "instructions");
    ws2 = new WebSocket(wsuri, "relay_backend");
    ws1.onmessage = function(e){doinstructions(e.data)};
    ...
    } // load
```
"""
function subprotocol # Methods added in include files
end

"""
`origin(request) => String`
Convenience function for checking which server / port address
the client claims its code was downloaded from.
The resource path can be found with target(req).
E.g.
```julia
function gatekeeper(req, ws)
    orig = WebSockets.origin(req)
        if startswith(orig, "http://localhost") || startswith(orig, "http://127.0.0.1")
            handlewebsocket(ws)
        end
    end
end
```
"""
function origin # Methods added in include files
end

"""
`writeguarded(websocket, message) => Bool`

Return true if write is successful, false if not.
The peer can potentially disconnect at any time, but no matter the
cause you will usually just want to exit your websocket handling function
when you can't write to it.

To check the errors (if you get any), temporarily set loging min_level to Logging.debug, e.g:

```julia
using WebSockets, Logging
global_logger(WebSocketLogger(stderr, Logging.Debug));
```
"""
function writeguarded(ws, msg)
    try
        write(ws, msg)
    catch err
        @debug err
        return false
    end
    true
end

"""
`readguarded(websocket) => (Vector, Bool)`

Return (data::Vector, true)
        or
        (Vector{UInt8}(), false)

The peer can potentially disconnect at any time, but no matter the
cause you will usually just want to exit your websocket handling function
when you can't write to it.

E.g.
```julia
while true
    data, success = readguarded(websocket)
    !success && break
    println(String(data))
end
```

To check the errors (if you get any), temporarily set loging min_level to Logging.debug, e.g:

```julia
using WebSockets, Logging
global_logger(WebSocketLogger(stderr, Logging.Debug));
```

"""
function readguarded(ws)
    data = Vector{UInt8}()
    success = true
    try
        data = read(ws)
    catch err
        @debug err
        data = Vector{UInt8}()
        success = false
    finally
        return data, success
    end
end

# import Base.show and add methods in this file
include("show_ws.jl")
end # module WebSockets


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

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 02329bbc-15a4-11ec-2522-df175bfb53e0
begin
    import Pkg
    # careful: this is _not_ a reproducible environment
    # activate the global environment
    Pkg.activate()

    using PlutoUI
	using FlightGNC
	using ComponentArrays, Transducers
	using Plots
end

# ╔═╡ 50becda1-a0e6-4e4c-98ca-a72992a5fc42
include("main_missiles.jl")

# ╔═╡ b9c304e7-9034-44d4-a4a6-a5d34d0baf58
begin
	include("sim_missiles.jl")

	# To avoid redefinition of a named (or an anonymous) function at each time tuning parameters and running simulation,
	# place the custom definition of a bias command function in a different cell from the cell containing sim_plot.

	custom_bias(s,x,t) = zeros(3)                           # option 1
	# function custom_bias(s,x,t)  return zeros(3) end      # option 2
	# custom_bias = (s,x,t) -> zeros(3)                     # option 3
end

# ╔═╡ b4aa1cc3-07de-4717-a39e-8a9c4b1c8512
md"""
# Missile Guidance Simulation
2D/3D Vector-Form Biased Proportional Navigation Guidance Law is implemented.
The bias command currently implemented considers target being stationary.

## Vector BPNG Equations
* Three-dimensional lead angle
$\sigma = \cos^{-1}\left(\hat{\mathbf{r}} \cdot \hat{\mathbf{v}}\right) \in \left[0,\pi\right]$

* Manoeuvre plane unit normal vector
$\hat{\mathbf{k}} = \frac{\mathbf{r} \times \mathbf{v}}{\left\|\mathbf{r} \times \mathbf{v}\right\|}$

* Predicted final velocity direction vector
$\begin{aligned}
	\hat{\mathbf{v}}_{f_{pred}} &= \hat{\mathbf{v}} \cos\left(\frac{N}{N-1}\sigma\right) - \left(\hat{\mathbf{k}}\times\hat{\mathbf{v}}\right)\sin\left(\frac{N}{N-1}\sigma\right)\\
	&= \hat{\mathbf{r}} \cos\left(\frac{1}{N-1}\sigma\right) - \left(\hat{\mathbf{k}}\times\hat{\mathbf{r}}\right)\sin\left(\frac{1}{N-1}\sigma\right)
\end{aligned}$

* Predicted final velocity direction error angle
$e = \cos^{-1}\left(\hat{\mathbf{v}}_{f_{pred}} \cdot \hat{\mathbf{v}}_{f_{d}}\right) \in \left[0,\pi\right]$

* Desired angular velocity of predicted final velocity direction vector
$\boldsymbol{\omega}_{f_{d}} = -K_{e}\left(r\right) \dot{r} e   \frac{\hat{\mathbf{v}}_{f_{pred}} \times \hat{\mathbf{v}}_{f_{d}}}{\left\|\hat{\mathbf{v}}_{f_{pred}} \times \hat{\mathbf{v}}_{f_{d}}\right\|} + K_{roll}\left(r\right) \hat{\mathbf{v}}_{f_{d}}$

* Bias angular velocity command
$\boldsymbol{\omega}_{bias} = K_{\sigma}\left(\sigma\right) \left[-\left(N-1\right) \left(\boldsymbol{\omega}_{f_{d}}\cdot \hat{\mathbf{k}}\right) \hat{\mathbf{k}}  + \sin\sigma \,\boldsymbol{\omega}_{f_{d}} \cdot \left\{\hat{\mathbf{r}} + \cot\left(\frac{1}{N-1}\sigma\right)\left(\hat{\mathbf{k}}\times\hat{\mathbf{r}}\right)\right\} \left(\hat{\mathbf{v}} \times \hat{\mathbf{k}}\right)\right]$

* LOS angular velocity
$\boldsymbol{\omega}_{r}^{r\perp} = \frac{\mathbf{v} \times \mathbf{r}}{\left\| \mathbf{r} \right\|^{2}} = -\frac{v\sin\sigma}{r}\hat{\mathbf{k}}$

* Lateral acceleration command
$\mathbf{a}_{BPNG}^{v\perp} = \left(N\boldsymbol{\omega}_{r}^{r\perp} + \boldsymbol{\omega}_{bias}\right) \times \mathbf{v}$

"""

# ╔═╡ 9fcf0c9b-d559-402f-883c-c53f61815c94
begin
	md"""
	 $\alpha$ = $(@bind α Slider(0.01:0.01:1; default=1, show_value=true))

	 $n$ = $(@bind n Slider(0 : 0.1 : 10; default=1, show_value=true))

	 $\gamma_{f_{d}}$ [deg] = $(@bind γ_f_d_deg Slider(-90:1:90; default=0, show_value=true))

	 $\chi_{f_{d}}$ [deg] = $(@bind χ_f_d_deg Slider(0:1:180; default=180, show_value=true))

	 $\sigma_{M_{\lim}}$ [deg] = $(@bind σ_M_lim_deg Slider(0:1:90; default=60, show_value=true))
	"""
end

# ╔═╡ b75636f4-b7f2-4356-9e2a-d90942c0071e
begin
	s_BPNG.Bias    = Bias_IACG_StationaryTarget
	s_BPNG.s_Bias  = ComponentArray(α = 1, δ = 0.01, n = n, r_ref = 10E3, k = 9, m = 10)
	s_BPNG.σ_M_lim = deg2rad(σ_M_lim_deg)
	γ_f_d = deg2rad(γ_f_d_deg)
	χ_f_d = deg2rad(χ_f_d_deg)
	s_BPNG.v̂_f_d = [cos(γ_f_d)*sin(χ_f_d); cos(γ_f_d)*cos(χ_f_d); sin(γ_f_d)]
	df, f = sim_plot(p_M_0, v_M_0, p_T_0, v_T_0, s_BPNG, A_M_max)
	f
end

# ╔═╡ bce435b0-c6d4-4ac7-93aa-727ed08c73e5
begin
	t = df.time
	r = df.sol |> Map(datum -> datum.r) |> collect
	plot(t,r, label = :false)
	annotate!((0.9,0.9), "min r = "*string(minimum(r)) )
end

# ╔═╡ Cell order:
# ╟─b4aa1cc3-07de-4717-a39e-8a9c4b1c8512
# ╟─02329bbc-15a4-11ec-2522-df175bfb53e0
# ╠═50becda1-a0e6-4e4c-98ca-a72992a5fc42
# ╠═b9c304e7-9034-44d4-a4a6-a5d34d0baf58
# ╟─9fcf0c9b-d559-402f-883c-c53f61815c94
# ╠═b75636f4-b7f2-4356-9e2a-d90942c0071e
# ╟─bce435b0-c6d4-4ac7-93aa-727ed08c73e5


text:
include("nodes.jl")
include("resolver.jl")

struct ComposerError
    context::Union{String, Nothing}
    context_mark::Union{Mark, Nothing}
    problem::Union{String, Nothing}
    problem_mark::Union{Mark, Nothing}
    note::Union{String, Nothing}

    function ComposerError(context=nothing, context_mark=nothing,
                           problem=nothing, problem_mark=nothing,
                           note=nothing)
        new(context, context_mark, problem, problem_mark, note)
    end
end

function show(io::IO, error::ComposerError)
    if error.context != nothing
        print(io, error.context, " at ", error.context_mark, ": ")
    end
    print(io, error.problem, " at ", error.problem_mark)
end

mutable struct Composer
    input::EventStream
    anchors::Dict{String, Node}
    resolver::Resolver
end

function compose(events)
    composer = Composer(events, Dict{String, Node}(), Resolver())
    @assert typeof(forward!(composer.input)) == StreamStartEvent
    node = compose_document(composer)
    if typeof(peek(composer.input)) == StreamEndEvent
        forward!(composer.input)
    else
        @assert typeof(peek(composer.input)) == DocumentStartEvent
    end
    node
end

function compose_document(composer::Composer)
    @assert typeof(forward!(composer.input)) == DocumentStartEvent
    node = compose_node(composer)
    @assert typeof(forward!(composer.input)) == DocumentEndEvent
    empty!(composer.anchors)
    node
end

function handle_event(event::AliasEvent, composer)
    anchor = event.anchor
    forward!(composer.input)
    haskey(composer.anchors, anchor) || throw(ComposerError(
            nothing, nothing, "found undefined alias '$(anchor)'", event.start_mark))
    return composer.anchors[anchor]
end

handle_error(event, composer, anchor) =
    anchor !== nothing && haskey(composer.anchors, anchor) && throw(ComposerError(
                "found duplicate anchor '$(anchor)'; first occurance",
                composer.anchors[anchor].start_mark, "second occurence",
                event.start_mark))

function handle_event(event::ScalarEvent, composer)
    anchor = event.anchor
    handle_error(event, composer, anchor)
    compose_scalar_node(composer, anchor)
end

function handle_event(event::SequenceStartEvent, composer)
    anchor = event.anchor
    handle_error(event, composer, anchor)
    compose_sequence_node(composer, anchor)
end

function handle_event(event::MappingStartEvent, composer)
    anchor = event.anchor
    handle_error(event, composer, anchor)
    compose_mapping_node(composer, anchor)
end

handle_event(event, composer) = nothing

function compose_node(composer::Composer)
    event = peek(composer.input)
    handle_event(event, composer)
end

function _compose_scalar_node(event::ScalarEvent, composer::Composer, anchor::Union{String, Nothing})
    tag = event.tag
    if tag === nothing || tag == "!"
        tag = resolve(composer.resolver, ScalarNode,
                      event.value, event.implicit)
    end

    node = ScalarNode(tag, event.value, event.start_mark, event.end_mark,
                      event.style)
    if anchor !== nothing
        composer.anchors[anchor] = node
    end

    node
end

compose_scalar_node(composer::Composer, anchor::Union{String, Nothing}) =
    _compose_scalar_node(forward!(composer.input), composer, anchor)

__compose_sequence_node(event::SequenceEndEvent, composer, node) = false
function __compose_sequence_node(event::Event, composer, node)
    push!(node.value, compose_node(composer))
    true
end

function _compose_sequence_node(start_event::SequenceStartEvent, composer, anchor)
    tag = start_event.tag
    if tag === nothing || tag == "!"
        tag = resolve(composer.resolver, SequenceNode,
                      nothing, start_event.implicit)
    end

    node = SequenceNode(tag, Any[], start_event.start_mark, nothing,
                        start_event.flow_style)
    if anchor !== nothing
        composer.anchors[anchor] = node
    end

    while (event = peek(composer.input)) !== nothing
        __compose_sequence_node(event, composer, node) || break
    end

    end_event = forward!(composer.input)
    node.end_mark = end_event.end_mark

    node
end

compose_sequence_node(composer::Composer, anchor::Union{String, Nothing}) =
    _compose_sequence_node(forward!(composer.input), composer, anchor)

__compose_mapping_node(event::MappingEndEvent, composer, node) = false
function __compose_mapping_node(event::Event, composer, node)
    item_key = compose_node(composer)
    item_value = compose_node(composer)
    push!(node.value, (item_key, item_value))
    true
end

function _compose_mapping_node(start_event::MappingStartEvent, composer::Composer, anchor::Union{String, Nothing})
    tag = start_event.tag
    if tag === nothing || tag == "!"
        tag = resolve(composer.resolver, MappingNode,
                      nothing, start_event.implicit)
    end

    node = MappingNode(tag, Any[], start_event.start_mark, nothing,
                       start_event.flow_style)
    if anchor !== nothing
        composer.anchors[anchor] = node
    end

    while true
        event = peek(composer.input)
        __compose_mapping_node(event, composer, node) || break
    end

    end_event = forward!(composer.input)
    node.end_mark = end_event.end_mark

    node
end

compose_mapping_node(composer::Composer, anchor::Union{String, Nothing}) =
    _compose_mapping_node(forward!(composer.input), composer, anchor)


text:
brlstring(brl::Number) = ':' * string(brl)
brlstring(::Nothing) = ""

"""
This function exists so that the same code in `create_newick` returns a string or prints to an IO (which is much faster), depending on the first argument
"""
_nwkout(o::String, s...) = string(o, s...)
function _nwkout(o::IO, s)
    print(o, s...)

    return o
end

function create_newick(o, p::RNode; brlengths=true)
    nc = outdegree(p)
    if nc > 0
        o = _nwkout(o, '(')
        for (i, c) in enumerate(children(p))
            o = create_newick(o, c; brlengths=brlengths)
            o = i < nc ? _nwkout(o, ',') : o
        end
        o = _nwkout(o, ')')
    end
    o = _nwkout(o, getlabel(p))
    o = brlengths ? _nwkout(o, brlstring(brlength(p))) : o

    return o
end

"""
    newick(node; brlengths=true, fmt=nothing)

Return a Newick string representation of the clade under a `node`. Will not include branch lengths if `brlengths` is set to `false`.
"""
newick(p::RNode; brlengths=true) = create_newick("", p; brlengths=brlengths) * ";"

newick(tree::AbstractTree; brlengths=true) = newick(getroot(tree); brlengths=brlengths)

"""
"""
function print_newick(io::IO, tree::AbstractTree; brlengths=true)
    create_newick(io, getnode(tree, getroot(tree)), brlengths=brlengths)
    print(io, ';')

    return nothing
end

max_stars_repo_name: eascarrunz/RootedTrees

max_issues_repo_name: eascarrunz/RootedTrees

max_forks_repo_name: eascarrunz/RootedTrees


text:
@testset "Genera" begin

  # TODO: move the remaining examples in the corresponding files in ~/test/QuadForm/Quad
  # Representatives

  Qx, x = FlintQQ["x"]
  K, a = NumberField(x^2 - 15)
  gens = [[a-1, 2*a-12, 0], [-471//70*a+881//14, 12*a-471, 1//70*a+39//14], [-7367*a+33891, 38904*a-212340, -194*a+1164], [2858191//5*a-1701731, -3700688*a+8438412, 103014//5*a-40352]]
  D = matrix(K, 3, 3, [38*a+150, 0, 0, 0, 2*a+8, 0, 0, 0, 302*a+1170])
  L = @inferred quadratic_lattice(K, gens, gram = D)
  @test length(@inferred genus_representatives(L)) == 1

  K, a = NumberField(x - 1)
  gens = [[-768, -1024, 0], [-768, -928, 160], [256, 320, -64]]
  D = matrix(K, 3, 3, [2, 0, 0, 0, 2, 0, 0, 0, 2])
  L = quadratic_lattice(K, gens, gram = D)
  @test length(genus_representatives(L)) == 2

  K, a = NumberField(x - 1)
  gens = [[-128, -1024, 0], [0, -2016, -32], [0, -512, 0]]
  D = matrix(K, 3, 3, [2, 0, 0, 0, 2, 0, 0, 0, 2])
  L = quadratic_lattice(K, gens, gram = D)
  @test length(genus_representatives(L)) == 2

  K, a = NumberField(x^2 - 2)
  gens = [[32*a-48, 32*a-64, 0], [0, 208*a+272, -24], [0, 512*a+768, -40*a-8]]
  D = matrix(K, 3, 3, [2, 0, 0, 0, 2, 0, 0, 0, 2])
  L = quadratic_lattice(K, gens, gram = D)
  @test length(genus_representatives(L)) == 2

  K, a = NumberField(x^3-x^2-2*x+1)
  gens = [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 1//2, 0, 1//2]]
  D = matrix(K, 4, 4, [1, 0, 0, 0, 0, 2, 0, 0, 0, 0, 1, 0, 0, 0, 0, 2])
  L = quadratic_lattice(K, gens, gram = D)
  @test length(genus_representatives(L)) == 2

  K, a = NumberField(x - 1)
  OK = maximal_order(K)
  p = prime_decomposition(OK, 2)[1][1]
  @test length(Hecke.local_genera_quadratic(K, p, rank = 2, det_val = 1)) == 8

  # Addition of genera

  Qx, x = PolynomialRing(FlintQQ, "x", cached = false)
  f = x - 1;
  K, a = number_field(f)
  D = matrix(K, 3, 3, [2, 0, 0, 0, -1, 0, 0, 0, -36]);
  gens = [[1, 0, 0], [1, 0, 0], [0, 1, 0], [0, 1, 0], [1//2, 1//2, 1//4], [1//2, 1//2, 1//4]]
  L = quadratic_lattice(K, gens, gram = D)
  p = prime_decomposition(maximal_order(K), 3)[1][1]
  fl, LL = @inferred Hecke.ismaximal_integral(L, p)
  @test !fl
  fl, _ = Hecke.ismaximal_integral(LL, p)
  @test fl

  Qx, x = PolynomialRing(FlintQQ, "x", cached = false)
  f = x - 1;
  K, a = number_field(f)
  D = matrix(K, 3, 3, [2, 0, 0, 0, -1, 0, 0, 0, -36]);
  V = quadratic_space(K, D)

  W = quadratic_space(K, diagonal_matrix([K(2), K(-36)]))

  @test @inferred Hecke.isrepresented_by(W, V)

  W = quadratic_space(K, diagonal_matrix([K(2), K(36)]))

  @test @inferred !Hecke.isrepresented_by(W, V)

end


text:
using MAT
using Random
using StatsBase
using SparseArrays
include("../include/HyperLocal/HyperLocal.jl")
include("../include/HyperLocal/SparseCardHyperLocal.jl")
include("../src/hypergraph-cond-utils.jl")
## Read in the matrix, if it isn't read in already

s = time()
@time M = matread(homedir()*"/Google Drive/Data Temp/stackoverflow_answer_H.mat")
LabelMatrix = M["LabelMatrix"]
LabelNames = M["LabelNames"]
MainLabels = M["MainLabels"]
H = M["H"]
order = round.(Int64,vec(sum(H,dims=2)))
d = vec(sum(H,dims=1))
volA = sum(d)
m,n = size(H)
Ht = sparse(H')
toc = time()-s
Edges = incidence2elist(H)
println("Done loading things into memory in $toc seconds.")

## Clusters: follows approach of Veldt et al. [KDD 2020]
Tconds = Vector{Float64}()
Tsizes = Vector{Int64}()
labels = Vector{Int64}()
for lab = 1:length(MainLabels)
    label = MainLabels[lab]
    T = findnz(LabelMatrix[:,label])[1]
    nT = length(T)
    condT, volT, cutT = tl_cond(H,T,d,1.0,volA,order)
    if condT < .2
        println("$nT \t $condT \t"*LabelNames[label])
        push!(Tconds,condT)
        push!(labels,label)
        push!(Tsizes,nT)
    end
end

## Save some seed sets
lnum = length(labels)
R_stats = zeros(lnum,3)
seed_sets = spzeros(n,lnum)
for lab = 1:length(labels)
    label = labels[lab]
    T = findnz(LabelMatrix[:,label])[1]
    nT = length(T)
    println("$lab \t $nT \t $label")

    # Generate a new seed set
    seednum = round(Int64,nT/20)
    grownum = round(Int64,min(nT*2,n/2-seednum))
    p = randperm(nT)
    Rstart = T[p[1:seednum]]
    OneHop = get_immediate_neighbors(H,Ht,Rstart)
    Rmore = BestNeighbors(H,d,Rstart,OneHop,grownum)
    R = union(Rmore,Rstart)
    Rs = findall(x->in(x,Rstart),R)     # Force seed nodes to be in output set
    prr, rer, f1r = PRF(T,R)
    R_stats[lab,:] = [prr, rer, f1r]
    seed_sets[Rstart,lab] .= 1
end
matwrite("Seed_sets_stack45.mat",Dict("R_stats"=>R_stats,"seed_sets"=>seed_sets))

max_stars_repo_name: nveldt/SparseCardDSFM

max_issues_repo_name: nveldt/SparseCardDSFM

max_forks_repo_name: nveldt/SparseCardDSFM


text:
using SimpleTraits.BaseTraits, SimpleTraits

export Add

"""
    Add{N} <: CompositeBlock{N}
    Add{N}(iterable) -> Add
    Add(blocks::AbstractBlock{N}...) -> Add
"""
struct Add{N} <: CompositeBlock{N}
    list::Vector{AbstractBlock{N}}

    Add{N}(list::Vector{AbstractBlock{N}}) where {N} = new{N}(list)
    Add{N}() where {N} = new{N}(AbstractBlock{N}[])
end

#Add{N}() where N = Add{N}(AbstractBlock{N}[])
Add{N}(blocks) where {N} = Add{N}(collect(AbstractBlock{N}, blocks))
Add(blocks::AbstractBlock{N}...) where {N} = Add{N}(blocks)

function mat(::Type{T}, x::Add{N}) where {N,T}
    length(x.list) == 0 && return Diagonal(zeros(T, 1 << N))
    mapreduce(x -> mat(T, x), +, x.list)
end

chsubblocks(x::Add{N}, it) where {N} = Add{N}(it)

function _apply!(r::AbstractRegister, x::Add)
    isempty(x.list) && return r
    length(x.list) == 1 && return _apply!(r, x.list[])

    res = mapreduce(blk -> _apply!(copy(r), blk), regadd!, x.list[1:end-1])
    _apply!(r, x.list[end])
    regadd!(r, res)
    r
end

export Add

subblocks(x::Add) = x.list
cache_key(x::Add) = map(cache_key, x.list)
Base.copy(x::Add{N}) where {N} = Add{N}(copy(x.list))
Base.similar(c::Add{N}) where {N} = Add{N}(empty!(similar(c.list)))

function Base.:(==)(lhs::Add{N}, rhs::Add{N}) where {N}
    (length(lhs.list) == length(rhs.list)) && all(lhs.list .== rhs.list)
end

for FUNC in [:length, :iterate, :getindex, :eltype, :eachindex, :popfirst!, :lastindex]
    @eval Base.$FUNC(x::Add, args...) = $FUNC(subblocks(x), args...)
end

Base.getindex(c::Add{N}, index::Union{UnitRange,Vector}) where {N} = Add{N}(getindex(c.list, index))
Base.setindex!(c::Add{N}, val::AbstractBlock{N}, index::Integer) where {N} =
    (setindex!(c.list, val, index); c)
Base.insert!(c::Add{N}, index::Integer, val::AbstractBlock{N}) where {N} =
    (insert!(c.list, index, val); c)
Base.adjoint(blk::Add{N}) where {N} = Add{N}(map(adjoint, subblocks(blk)))

## Iterate contained blocks
occupied_locs(c::Add) = (unique(Iterators.flatten(occupied_locs(b) for b in subblocks(c)))...,)

# Additional Methods for Add
Base.push!(c::Add{N}, val::AbstractBlock{N}) where {N} = (push!(c.list, val); c)

function Base.push!(c::Add{N}, val::Function) where {N}
    push!(c, val(N))
end

function Base.append!(c::Add, list)
    for blk in list
        push!(c, blk)
    end
    c
end

function Base.prepend!(c::Add, list)
    for blk in list[end:-1:1]
        insert!(c, 1, blk)
    end
    c
end

YaoBase.ishermitian(ad::Add) = all(ishermitian, ad.list) || ishermitian(mat(ad))

max_stars_repo_name: QuantumBFS/YaoStdCircuit

max_issues_repo_name: QuantumBFS/YaoStdCircuit


text:
using Parameters, StaticArrays, CSV, DataFrames

function runfig2()
    p = Params(Ln = 1440, Ls = 0, scale = 40, U = 1e-7, λ = 10, α = 2, EZ = SA[0, 0, 0],
        Δ = 0., μN = 0.);
    p = reconstruct(p, EZ = SA[0, 0.001, 0])
    p_EZ = reconstruct(p, EZ = SA[0, 2.8, 0])
    mat = fig2bands([p, p, p_EZ, p_EZ],
       ["armchair", "zigzag", "armchair", "zigzag"])
    writecsv(mat, [p, p, p_EZ, p_EZ], "fig2", ["1", "2", "3", "4"])
end

fig2bands(listp::Array, edgelist) =
    [fig2bands(listp[i], edgelist[i]) for i in 1:length(listp)]

function fig2bands(p, which)
    kpoints = 301
    numbands = 56
    if which == "armchair"
        kmin, kmax = (-0.15, 0.15)
        axis = (0, 1)
        h = nanoribbonN(p, axis);
        b = bandstructure(h,
            cuboid((kmin, kmax); subticks = kpoints),
            mapping = x -> 2pi*x,
            method = ArpackPackage(nev=numbands, maxiter = 300, sigma=-0.00001));
    else
        kmin, kmax = (0.2, 0.8)
        axis = (1, 0)
        h = nanoribbonN(p, axis);
        b = bandstructure(h, cuboid((kmin, kmax), subticks = kpoints),
            mapping = x -> 2pi*x,
            method = ArpackPackage(nev=numbands, maxiter = 300, sigma=-0.00001))

    end
    karray = collect(range(kmin, kmax, length = kpoints))
    ϵs = zeros(ComplexF64, numbands, kpoints)
    for i in 1:length(karray)
        ϵs[:,i], _ = b.diag((karray[i],))
    end
    return karray, ϵs, b
end

max_stars_repo_name: fernandopenaranda/MBSinBLG

max_issues_repo_name: fernandopenaranda/MBSinBLG

max_forks_repo_name: fernandopenaranda/MBSinBLG


text:
immutable FilterContext
    time_step::Int64
end

function forward!(filters::Vector{Filter}, l::Layer; time_step=0)
    context = FilterContext(time_step)
    for i in 1:length(filters)
        if filters[i].enabled
            forward!(filters[i], l, context)
        end
    end
end

function backprop!(filters::Vector{Filter}, l::Layer; time_step=0)
    context = FilterContext(time_step)
    for i in 1:length(filters)
        if filters[i].enabled
            backprop!(filters[i], l, context)
        end
    end
end

function forward!(f::Filter, l::Layer, context::FilterContext)
    # No-op
end

function backprop!(f::Filter, l::Layer, context::FilterContext)
    # No-op
end

# ----------------------------------------------------------------
immutable Dropout <: Filter
    enabled::Bool
    drop::Float32
    scale::Float32
    mask::Signal
    Dropout(drop) = new(true, drop, 1.0/(1.0 - drop))
end

function forward!(drop::Dropout, l::Layer, context::FilterContext)
    mask = (rand(size(l.Z)) .> drop.drop)
    copy!(drop.mask, mask)
    drop!(drop, l.Z)
end

function backprop!(drop::Dropout, l::Layer, context::FilterContext)
    drop!(drop, l.ΔE)
end

function forward!(drop::Dropout, l::RecurrentLayer, context::FilterContext)
    error("Not yet implemented")
end

function backprop!(drop::Dropout, l::RecurrentLayer, context::FilterContext)
    error("Not yet implemented")
end

function drop!(drop::Dropout, X::Signal)
    for col=1:size(X,2), row=1:size(X,1)
        X[row, col] = X[row, col] * drop.scale * drop.mask[row, col]
    end
end

# ----------------------------------------------------------------
immutable GradientClip <: Filter
    enabled::Bool
    threshold::Float32
    GradientClip(threshold) = new(true, threshold)
end

function backprop!(grad::GradientClip, l::Layer, context::FilterContext)
    gradient_clip!(l.ΔE, grad.threshold)
end

function backprop!(grad::GradientClip, l::RecurrentLayer, context::FilterContext)
    t = context.time_step
    gradient_clip!(sub(l.ΔE,:,t:t),  grad.threshold)
end

function gradient_clip!(g, limit::Float32)
    nm = norm(g,1)
    if nm > limit
        rate = limit / nm
        ng = @in1! g .* rate
    end
    g
end


text:
struct LinearAnalysis
    options::Dict{Symbol, Any}
end

LinearAnalysis(; kwargs...) = LinearAnalysis(Dict{Symbol, Any}(kwargs))

add_intercept_column(x::AbstractVector{T}) where {T} = [ones(T, length(x)) x]

# TODO: add multidimensional version
function (l::LinearAnalysis)(le::Entry)
    return splitapply(le) do entry
        labels, mappings = map(getlabel, entry.mappings), map(getvalue, entry.mappings)
        x, y = mappings.positional
        wts = get(mappings, :wts, similar(x, 0))
        npoints = get(l.options, :npoints, 200)
        interval = get(l.options, :interval, length(wts) > 0 ? nothing : :confidence)
        dropcollinear = false
        lin_model = try
            GLM.lm(add_intercept_column(x), collect(y); wts, dropcollinear)
        catch e
            @warn "linear fit not possible"
            nothing
        end
        isnothing(lin_model) && return Entry[]
        x̂ = range(extrema(x)..., length=npoints)
        pred = GLM.predict(lin_model, add_intercept_column(x̂); interval)
        if !isnothing(interval)
            ŷ, lower, upper = map(vec, pred) # GLM prediction returns matrices
            default_plottype = LinesFill
            labeled_result = map(Labeled, vcat(labels.positional, ["", ""]), [x̂, ŷ, lower, upper])
        else
            ŷ = vec(pred) # GLM prediction returns matrix
            default_plottype = Lines
            labeled_result = map(Labeled, labels.positional, [x̂, ŷ])
        end
        return Entry(
            AbstractPlotting.plottype(entry.plottype, default_plottype),
            Arguments(labeled_result),
            entry.attributes
        )
    end
end

"""
    linear(; interval)

Compute a linear fit of `y ~ 1 + x`. An optional named mapping `wts` determines the weights.
Use `interval` to specify what type of interval the shaded band should represent.
Valid values of interval are `:confidence` delimiting the uncertainty of the predicted
relationship, and `:prediction` delimiting estimated bounds for new data points.
"""
linear(; kwargs...) = Layer((LinearAnalysis(; kwargs...),))


text:
using DesignDriveCoilFilt
using Test

@testset "DesignDriveCoilFilt.jl" begin
    # Write your tests here.
end


text:
plot(
    p0 ->
        FRGRealTime.VImintqs(
            p0,
            0.0001,
            80.0,
            145.0,
            4.0,
            1.0,
            800.0,
            m2fun,
            lamfun,
        ) + FRGRealTime.VImintqs_delta1(
            p0,
            0.0001,
            80.0,
            145.0,
            4.0,
            1.0,
            800.0,
            m2fun,
            lamfun,
        ),
    1.0,
    800.0,
    label = "I2, p0-q0",
    legend = :topleft,
)

quadgk(k->FRGRealTime.dkV4piResimple(
    0.01,
    0.0001,
    Epi(100.0, m2fun(100.0)),
    100.0,
    0.0,
    k,
    m2fun(k),
    145.0,
    lamfun(k),
    4.0,
),800.0,100.0)[1]

plot(qs->quadgk(k->FRGRealTime.dkV4piResimple(
    0.000001,
    0.0001,
    0.0001,
    qs,
    0.0,
    k,
    m2fun(k),
    145.0,
    lamfun(k),
    4.0,
),800.0,100.0)[1]-64,1.0,100.0)

plot(
    k -> FRGRealTime.dkF1Tilde3(k,m2fun(k), 145.0),
    2.0,
    100.0,
)

plot(
    k -> central_fdm(5, 1)(
        qsmax ->
            FRGRealTime.PvdkF1Tildeps(0.001, qsmax, k, m2fun(k), 145.0),
        2.0,
    ),
    0.01,
    10.0,
)

v1 = collect(1.0:0.2:800.0)

v2 = exp.(collect(log(1e-8):(log(1e-4)-log(1e-8))/1500:log(1e-4)))

outv1 = map(v1) do p0
    FRGRealTime.VImintqs(
        p0,
        0.0001,
        k_temp,
        145.0,
        4.0,
        1.0,
        800.0,
        m2fun,
        lamfun,
    ) + FRGRealTime.VImintqs_delta1(
        p0,
        0.0001,
        k_temp,
        145.0,
        4.0,
        1.0,
        800.0,
        m2fun,
        lamfun,
    )
end

FRGRealTime.loopfunppfix.(v2, 1e-6, 2.0, 2.0, 10.0)

plot(
    p0 -> FRGRealTime.loopfunppfix(p0, 1e-6, 2.0, 2.0, 10.0),
    1e-8,
    1e-4,
    xaxis = :log,
)

writedlm("$dir/k=50.dat", 3 * outv1)

plot!(v1, outv1 / 100^3)
dir = "/home/tyy/Documents/CTP-fRG-Test/Real-Time-data/Im/Loop"
k_temp = 50
writedlm("$dir/looppp_m=2.dat", FRGRealTime.loopfunppfix.(v2, 1e-6, 2.0, 2.0, 10.0))

k_temp = [50, 80, 100, 200, 300, 400, 500, 600, 700]
writedlm("$dir/klist.dat", k_temp)
writedlm("$dir/q0list.dat", Epi.(k_temp, m2fun.(k_temp)))

Epi.(k_temp, m2fun.(k_temp))

Epi(200, m2fun(200))


text:
using DataFrames
using CSV
using Plots
#%%
data_jacobi = open("jacobi.dat")
df_jacobi=CSV.read(data_jacobi, header=false, delim=" ")
x_jacobi=df_jacobi[1]
y_jacobi=df_jacobi[2]
z_jacobi=reshape(df_jacobi[3],(32,32))

data_gs = open("gauss_seidel.dat")
df_gs=CSV.read(data_gs, header=false, delim=" ")
x_gs=df_gs[1]
y_gs=df_gs[2]
z_gs=reshape(df_gs[3],(32,32))

data_s = open("sor.dat")
df_s=CSV.read(data_s, header=false, delim=" ")
x_s=df_s[1]
y_s=df_s[2]
z_s=reshape(df_s[3],(32,32))

data_sol = open("sol.dat")
df_sol=CSV.read(data_sol, header=false, delim=" ")
x_sol=df_sol[1]
y_sol=df_sol[2]
z_sol=reshape(df_sol[3],(32,32))

x=y_sol[1:32]
y=x
#%%
plot(x,y,z_jacobi, fill=true, label="jacobi")
xlabel!("x")
ylabel!("y")
title!("jacobi")
plot!(dpi=600)
savefig("jacobi.png")
# xaxis!(:log10)
# yaxis!(:log10)
#%%

plot(x,y,z_gs, fill=true, label="gauss_seidel")
xlabel!("x")
ylabel!("y")
title!("gauss seidel")
plot!(dpi=600)
savefig("gs.png")

#%%

plot(x,y,z_s, fill=true, label="sor")
xlabel!("x")
ylabel!("y")
title!("sor")
plot!(dpi=600)
savefig("sor.png")

max_stars_repo_name: aurelio-amerio/corso-cpp


text:
#treat thsi table as a one-off since it has no common elements w/ other tables
function maketable8(dist::DataFrame;
    panela::Bool = true,
    panelb::Bool = true,
    panelc::Bool = true)

  ################common data
  #println("Size of file: $(length(unique(dist.eid)))")
  #println(describe(dist))
  #error("stop")

  ###################PANEL A############

  T8_PANEL_A_NAME::String = "t8panela-$(REPLICATION_TYPE[])-$(DIST_TYPE[])-$(OUT_SUFFIX[])"

  T8_PANEL_A_ROWKEYS::Vector{Symbol} = [
    :firstlast,
    :alldist,
    :ordinary,
    :match,
    :oneex,
    :dclrexok]

  T8_PANEL_A_ROWLABEL::Dict = Dict(
    :firstlast => "First and Last Dividends",
    :alldist => "All Distributions on CRSP",
    :ordinary => "Ordinary `CRSP 1232' Dividends",
    :match => "Matches to CRSP",
    :oneex => "Only 1 ex-date in window",
    :dclrexok => "\$\\ge$DCLR_EX_DAYS\$ Between Declaration and Ex-Date")

  panela && table8panela(dist,
      tablename= T8_PANEL_A_NAME,
      rowkeys=T8_PANEL_A_ROWKEYS,
      rowlabel=T8_PANEL_A_ROWLABEL,
      dclrexdays = DCLR_EX_DAYS,
      colnames = [[""]],
      widthcolnames = [[1]])

  ###################PANEL B############

  sdist::SubDataFrame = view(dist, dist.oneex .* dist.match .* dist.ordinary .* dist.inrange .* dist.primaryyearrange,:)

  T8_PANEL_B_ROWKEYS::Vector{Symbol} = [
    :dclrex,
    :cumulative,
    :Mcumulative,
    :after,
    :dclrex2,
    :cumulative2,
    :Mcumulative2,]

  T8_PANEL_B_ROWLABEL::Dict = Dict(
    :dclrex=>"Dclrtn to Ex",
    :cumulative=>"Cumulative",
    :Mcumulative=>"\$-\$Cumulative",
    :after=>"\\midrule After Selection",
    :dclrex2=>"\\midrule Dclrtn to Ex",
    :cumulative2=>"Cumulative",
    :Mcumulative2=>"\$-\$Cumulative",)

  T8_PANEL_B_COLKEYS::Vector{Int} = collect(0:9)
  T8_PANEL_B_COLNAMES::Vector{Vector{String}} = [["Days"; (string).(T8_PANEL_B_COLKEYS)]]

  T8_PANEL_B_NAME::String = "t8panelb-$(REPLICATION_TYPE[])-$(DIST_TYPE[])-$(OUT_SUFFIX[])"

  panelb && table8panelb(sdist,
    rowkeys=T8_PANEL_B_ROWKEYS,
    rowlabel=T8_PANEL_B_ROWLABEL,
    tablename=T8_PANEL_B_NAME,
    colkeys = T8_PANEL_B_COLKEYS,
    colnames = T8_PANEL_B_COLNAMES,
    dclrexdays = DCLR_EX_DAYS
  )

  ###################PANEL C############
  sdist = view(dist, dist.primary,:)
  T8_PANEL_C_ROWKEYS::Vector{Symbol} = [
    :return,
    :returnabs,
    :yield]

  T8_PANEL_C_ROWLABEL::Dict = Dict(
    :return=>"Net-of-Market Rates of\\\\Return, \$-12\$ to \$+12\$ days",
    :returnabs=>"\\dots in abs values",
    :yield=>"Div Yield \$\\delta, |\\ge 9\$ days")

  T8_PANEL_C_COLKEYS::Vector{Symbol} = [
    :min,
    :median,
    :max,
    :pos,
    :mean,
    :sd,
    :t,
    :N,
    :Nwinsorized]

  T8_PANEL_C_COLLABEL::Dict = Dict(
    :min=>"Min",
    :median=>"Median",
    :max=>"Max",
    :pos=>"\\%Pos",
    :mean=>"Mean",
    :sd=>"SD",
    :t=>"Trd-days",
    :N=>"\\#",
    :Nwinsorized=>"\\# Winsrzd")

  T8_PANEL_C_NAME::String = "t8panelc-$(REPLICATION_TYPE[])-$(DIST_TYPE[])-$(OUT_SUFFIX[])"

  panelc && table8panelc(sdist,
    rowkeys=T8_PANEL_C_ROWKEYS,
    rowlabel=T8_PANEL_C_ROWLABEL,
    tablename=T8_PANEL_C_NAME,
    colkeys = T8_PANEL_C_COLKEYS,
    collabel = T8_PANEL_C_COLLABEL,
    dclrexdays = DCLR_EX_DAYS
  )
end

function table8panelc(dist::AbstractDataFrame;
  rowkeys::Vector{Symbol} = error("rowkeys is required"),
  rowlabel::Dict = error("rowlabel is required"),
  colkeys::Vector{Symbol} = error("colkeys is required"),
  collabel::Dict = error("collabel is required"),
  tablename::String = error("table name is required"),
  dclrexdays::Int = error("dclrexdays is required"),
  tablepath::String = TABLE_PATH,
  Fret::Symbol = DIST_RET,
  decimals::Int = 3,)

  local distagg::DataFrame

  colnames::Vector{Vector{String}} = [[(k::Symbol->string(collabel[k])).(colkeys)...;]]

  Nrows::Int = length(rowkeys)
  Ncols::Int = length(colnames[1])

  #pre-allocate space and index
  descrownames::Vector{String} = (s::Symbol->rowlabel[s]).(rowkeys)
  desccontent::Vector{Vector{String}} = (i->Vector{String}(undef, Ncols)).(1:Nrows)

  #index the rows and cols for readability
  rowidx::Dict = Dict(s=>desccontent[i] for (i,s::Symbol) ∈ enumerate(rowkeys))
  colidx::Dict = Dict(s=>i for (i,s::Symbol) ∈ enumerate(colkeys))

  #does the formatting as required
  function plusnum(num::Real, d::Int=decimals)
    local s::String
    (abs(num) < 10^-10)  && (num = 0)
    if num > 10^-10
      s = "+$(num2str(num, d, Ints=true))"
    else
      s = "$(num2str(num, d, Ints=true))"
    end
    println("$s: $num")
    return s
  end

  #handles the common elements across rows
  function writepartialrow(r::Vector{String}, rets::Vector{Float64}, scaling::Int)
    r[colidx[:min]] = plusnum(minimum(rets) * scaling)
    r[colidx[:median]] = plusnum(median(rets) * scaling)
    r[colidx[:max]] = plusnum(maximum(rets) * scaling)
    r[colidx[:mean]] = num2str(mean(rets) * scaling, decimals)
    r[colidx[:sd]] = num2str(std(rets) * scaling, decimals)
  end

  rets::Vector{Float64} = collect(skipmissing(dist[!, Fret]))
  absrets::Vector{Float64} = (abs).(rets)

  #grab the dividend yield for each event
  sdist = view(dist, dist.dclrexok, :)
  distagg = aggregate(groupby(view(sdist, sdist.day.==0,
    [:eid, :yield, :winsorizedyield, :dclrexdays, :dclrextradingdays]), :eid), x->x[1])

  yields::Vector{Float64} = collect(skipmissing(distagg.yield_function))

  writepartialrow(rowidx[:return], rets, 100)
  writepartialrow(rowidx[:returnabs], absrets, 100)
  writepartialrow(rowidx[:yield], yields, 100)

  #now that we have done the common elements, compute the rest
  #start with return
  rowidx[:return][colidx[:pos]] = "$(plusnum(sum(rets .> 0)/length(rets) * 100,1))\\%"

  meandays::Float64 = mean(skipmissing(distagg.dclrextradingdays_function))
  rowidx[:return][colidx[:t]] = "$(plusnum(meandays, 1))"
  rowidx[:return][colidx[:N]] = "$(length(rets))"

  Nwinsorized::Int = sum(skipmissing(dist.winsorized))
  rowidx[:return][colidx[:Nwinsorized]] = "$Nwinsorized"

  #now compute additional rows for the divdend yield
  rowidx[:yield][colidx[:N]] = "$(length(yields))"

  Nwinsorizedyield::Int = sum(skipmissing(distagg.winsorizedyield_function))
  rowidx[:yield][colidx[:Nwinsorized]] = "$Nwinsorizedyield"

  #now compute the "stub" blanks
  rowidx[:returnabs][colidx[:pos]] = ""
  rowidx[:returnabs][colidx[:t]] = ""
  rowidx[:returnabs][colidx[:N]] = ""
  rowidx[:returnabs][colidx[:Nwinsorized]] = ""
  rowidx[:yield][colidx[:pos]] = ""
  rowidx[:yield][colidx[:t]] = ""

  tt::String = textable(colnames=colnames,
    descrownames=descrownames,
    desccontent=desccontent)

  writetextable(tt, path=tablepath, outname="$tablename.tex")

end

function table8panela(dist::DataFrame;
    tablename::String = error("tablename is required"),
    tablepath::String = TABLE_PATH,
    decimals::Int = DEFAULT_DECIMALS,
    colnames::Vector{Vector{String}} = error("colnames is required"),
    rowkeys::Vector{Symbol} = error("rowkeys is required"),
    rowlabel::Dict = error("rowlabel is required"),
    dclrexdays::Int = error("dclrexdays is required"),
    widthcolnames::Vector{Vector{Int}} =  error("widthcolnames is required"))

  local parallel::Bool = PARALLEL[] && (!DISABLE_FM_PARALLEL)
  local sdist::SubDataFrame = view(dist, dist.primaryyearrange, :)
  descrownames::Vector{String} = (s::Symbol->rowlabel[s]).(rowkeys)

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

  #preallocation and housekeeping
  Ncols::Int = sum(widthcolnames[1])
  Nrows::Int = length(descrownames)
  desccontent::Vector{Vector{String}} = (i->Vector{String}(undef, Ncols)).(1:Nrows)
  alignmentstring::String = "lc"

  #make the content
  minyr, minmo = yearmonth(minimum(dist.exdate))
  maxyr, maxmo = yearmonth(maximum(dist.exdate))

  #index the rows for readability
  rowindex::Dict = Dict(s=>desccontent[i] for (i,s::Symbol) ∈ enumerate(rowkeys))

  rowindex[:firstlast][1] = "\\text{$minyr/$minmo to $maxyr/$maxmo}"

  rowindex[:alldist][1] = "\\text{N=}$(length(unique(sdist.eid)))"

  sdist = view(sdist, sdist.ordinary, :)
  rowindex[:ordinary][1] = "\\text{N=}$(length(unique(sdist.eid)))"

  sdist = view(sdist, sdist.match, :)
  rowindex[:match][1] = "\\text{N=}$(length(unique(sdist.eid)))"

  sdist = view(sdist, sdist.oneex, :)
  rowindex[:oneex][1] = "\\text{N=}$(length(unique(sdist.eid)))"

  sdist = view(sdist, sdist.dclrexok, :)
  rowindex[:dclrexok][1] = "\\text{N=}$(length(unique(sdist.eid)))"

  #println(desccontent)

  #display(desccontent)
  #construct the tex table
  tt::String = textable(colnames=colnames,
    descrownames=descrownames,
    desccontent=desccontent,
    alignmentstring=alignmentstring,
    widthcolnames=widthcolnames)

  writetextable(tt, path=tablepath, outname="$tablename.tex")
end

function table8panelb(dist::AbstractDataFrame;
  rowkeys::Vector{Symbol} = error("rowkeys is required"),
  rowlabel::Dict = error("rowlabel is required"),
  colkeys::Vector{Int} = error("colkeys is required"),
  colnames::Vector{Vector{String}} = error("collabel is required"),
  tablename::String = error("table name is required"),
  dclrexdays::Int = error("dclrexdays is required"),
  tablepath::String = TABLE_PATH,
  decimals::Int = DEFAULT_DECIMALS,)

  Nrows::Int = length(rowkeys)
  Ncols::Int = length(colnames[1]) - 1

  #pre-allocate space and index
  descrownames::Vector{String} = (s::Symbol->rowlabel[s]).(rowkeys)
  desccontent::Vector{Vector{String}} = (i->Vector{String}(undef, Ncols)).(1:Nrows)
  desccondentidx::Dict = Dict(k=>desccontent[i] for (i::Int,k::Symbol) ∈ enumerate(rowkeys))

  #compute the number of events for each length of time between dclr and ex
  alldclrexdays::Vector{Int} = collect(0:(maximum(colkeys)+dclrexdays -1))
  dclrexdensity::Dict = Dict(Δ=>
    length(unique(view(dist, dist.dclrexdays .== Δ, :eid))) for Δ::Int ∈ alldclrexdays)

  #some important totals
  totevent::Int = length(unique(dist.eid))
  totselected::Int = length(unique(view(dist, dist.dclrexdays .≥ dclrexdays, :eid)))
  #construct the table
  for (i::Int,Δ::Int) ∈ enumerate(colkeys)

    dclrex::Int = dclrexdensity[Δ]
    desccondentidx[:dclrex][i] = num2str(dclrex, 0)

    #compute the cumulative sum
    cumulative::Int = sum((d->dclrexdensity[d]).(alldclrexdays[alldclrexdays .≤ Δ]))
    desccondentidx[:cumulative][i] = num2str(cumulative,0)

    #back out the negative cumulative sum
    Mcumulative::Int = totevent - cumulative
    desccondentidx[:Mcumulative][i] = num2str(Mcumulative,0)

    #second half table header
    Δ⁺::Int = Δ + 8
    desccondentidx[:after][i] = "$Δ⁺"

    if Δ⁺ == 8 #sepecial case to initialize the second part of the table
      desccondentidx[:dclrex2][i] = ""
      desccondentidx[:cumulative2][i] = ""
      desccondentidx[:Mcumulative2][i] = num2str(totselected,0)
    else
      dclrex2::Int = dclrexdensity[Δ⁺]
      desccondentidx[:dclrex2][i] = num2str(dclrex2, 0)

      #sums the density from the minimum selected period to Δ⁺
      cumulative2::Int = sum((d->
        dclrexdensity[d]).(alldclrexdays[(t->(t≥dclrexdays) && (t≤Δ⁺)).(alldclrexdays)]))

      desccondentidx[:cumulative2][i] = num2str(cumulative2,0)

      Mcumulative2::Int = totselected - cumulative2
      desccondentidx[:Mcumulative2][i] = num2str(Mcumulative2,0)
    end
  end

  #println(desccontent)

  tt::String = textable(colnames=colnames,
    descrownames=descrownames,
    desccontent=desccontent,
    rowlabelheader=true)

  writetextable(tt, path=tablepath, outname="$tablename.tex")

end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
struct KNN
  kcoeff::Float64
  neighbors::Data{Weighted}
  weights::Vector{Vector{Float64}}
end
get_bestk(n, knn::KNN) = round(Int, knn.kcoeff * sqrt(n))

"Find the best rule for k on the given data `d`."
function find_k_coeff(d::Data{Weighted}, k_range=1:100)
  M = length(d.y_T)

  best_k = zeros(Int, M)
  for m = 1:M
    X_m = d.X_T[m]
    y_m = d.y_T[m]
    n_m = length(y_m)

    max_k = min(maximum(k_range), n_m - 1)
    kdtree = NearestNeighbors.KDTree(Matrix(X_m'))
    # Find the k+1 nearest neighbors to each point, sorted by distance
    # Each point has itself as the first neighbor, which we will ignore
    idxs, _ = NearestNeighbors.knn(kdtree, Matrix(X_m'), max_k + 1, true)

    # Get the outcomes for each of the neighbors
    # Ignore the first neighbor since it's just the point itself
    cum_means = Vector{Vector{Float64}}(undef, n_m)
    for i = 1:n_m
      outcomes = y_m[idxs[i][2:end]]
      cum_means[i] = cumsum(outcomes) ./ (1:max_k)
    end

    # Find the k that minimizes the MAE of kNN predictions
    best_mae = Inf
    for k in k_range
      k >= n_m && continue
      mae = 0.0
      for i = 1:n_m
        mae += abs(cum_means[i][k] - y_m[i])
      end
      mae /= n_m

      if mae < best_mae
        best_mae = mae
        best_k[m] = k
      end
    end
  end

  sqrt_n = sqrt.(length.(d.y_T))
  path = GLMNet.glmnet(reshape(sqrt_n, M, 1), best_k)
  k_coeff = path.betas[1, end]
end

function getoutcomes(newdata_norm::Data{Normalized}, knn::KNN)
  M = length(knn.neighbors.X_T)

  n = size(newdata_norm.X, 1)
  outcomes = Matrix{Float64}(undef, n, M)

  for m = 1:M
    # Weight all points according to candidate treatment
    current_weighted = weight_data(newdata_norm, knn.weights, m)

    # Get all relevant weighted training points for candidate treatment
    X_m = knn.neighbors.X_T[m]
    y_m = knn.neighbors.y_T[m]

    k = get_bestk(length(y_m), knn)
    k = max(min(k, length(y_m)), 1)

    # Get k nearest neighbors for each new point from the neighbor pool
    # We don't need the neighbors in sorted order
    kdtree = NearestNeighbors.KDTree(Matrix(X_m'))
    idxs, _ = NearestNeighbors.knn(kdtree, Matrix(current_weighted.X'), k,
                                   false)

    # Get the mean outcome for each point from the neighbors
    for i = 1:n
      outcome = 0.0
      idx = idxs[i]
      for j = 1:k
        outcome += y_m[idx[j]]
      end
      outcomes[i, m] = outcome / k
    end
  end

  outcomes
end

"Trains kNN model on training data `d`."
function train(d::Data{Normalized}, ::Val{:knn})
  d_weighted, weights = weight_data(d)

  k_coeff = find_k_coeff(d_weighted)
  KNN(k_coeff, d_weighted, weights)
end

"Estimates counterfactuals on the dataset `d` using kNN."
function getcounterfactuals(d::Data{Normalized}, impute_method::Val{:knn})
  knn = train_knn(d, impute_method)
  getoutcomes(d, knn)
end


text:
struct SOS
    order::Int
    indices::Vector{Int}
    weights::Vector{Float64}
    SOS(order, indices, weights) = new(order, indices, weights)
end

mutable struct MPSModel <: NEOSModel
    solver::NEOSSolver
    xmlmodel::String
    last_results::String
    ncol::Int
    nrow::Int
    A::SparseMatrixCSC{Float64, Int}
    collb::Vector{Float64}
    colub::Vector{Float64}
    f::Vector{Float64}
    rowlb::Vector{Float64}
    rowub::Vector{Float64}
    sense::Symbol
    colcat::Vector{Symbol}
    sos::Vector{SOS}
    objVal::Float64
    reducedcosts::Vector{Float64}
    duals::Vector{Float64}
    solution::Vector{Float64}
    status::Symbol

    MPSModel(solver) = new(solver, "", "", 0, 0, sparse(Int[], Int[], Float64[]), Float64[], Float64[], Float64[], Float64[], Float64[], :Min, Symbol[], SOS[], 0.0, Float64[], Float64[], Float64[], NOTSOLVED)
end

MPB.LinearQuadraticModel(s::NEOSSolver{S,:MPS}) where {S} = MPSModel(s)

function neos_writexmlmodel!(m::MPSModel)
    io = IOBuffer()
    print(io, "<MPS>")
    mps_writer_sos = [(s.order, s.indices, s.weights) for s in m.sos]
    MPSWriter.write(io, m.A, m.collb, m.colub, m.f, m.rowlb, m.rowub,
                    m.sense, m.colcat, mps_writer_sos, Array{Float64}(undef,0,0))
    print(io, "</MPS>")
    # Convert the model to MPS and add
    m.xmlmodel = replace(m.solver.template, r"<MPS>.*</MPS>"is => String(take!(io)))
end

function MPB.loadproblem!(m::MPSModel, A, collb, colub, f, rowlb, rowub, sense)
    # @assert length(collb) == length(colub) == length(f)
    m.ncol = length(f)
    m.nrow = length(rowlb)
    m.A = sparse(A)
    m.collb = collb
    m.colub = colub
    m.colcat = fill(:Cont, m.ncol)
    m.f = f
    m.rowlb = rowlb
    m.rowub = rowub
    m.sense = sense
end

function anyints(m::MPSModel)
    for i in m.colcat
	if i == :Int || i==:Bin
	    return true
	end
    end
    return false
end

function parseresults!(m::MPSModel, job)
    m.status = SOVLERERROR
    parse_status!(m.solver, m)
    if m.status == OPTIMAL
	parse_objective!(m.solver, m)
	m.solution = zeros(m.ncol)
	parse_solution!(m.solver, m)
	if m.solver.provides_duals && length(m.sos) == 0 && !anyints(m)
	    parse_duals!(m.solver, m)
	end
	if m.sense == :Max
	    # Since MPS does not support Maximisation
	    m.objVal = -m.objVal
	end
    end
end

function addsos1!(m::MPSModel, indices::Vector, weights::Vector)
    if !m.solver.solves_sos
	throw(NEOSSolverError("Special Ordered Sets of type I are not supported by $(typeof(m.solver)). Try a different solver instead"))
    end
    push!(m.sos, SOS(1, indices, weights))
end
function addsos2!(m::MPSModel, indices::Vector, weights::Vector)
    if !m.solver.solves_sos
	throw(NEOSSolverError("Special Ordered Sets of type II are not supported by $(typeof(m.solver)). Try a different solver instead"))
    end
    push!(m.sos, SOS(2, indices, weights))
end

function MPB.status(m::MPSModel)
    m.status
end

function MPB.getobjval(m::MPSModel)
    m.objVal
end

function MPB.getsolution(m::MPSModel)
    m.solution
end

function MPB.getreducedcosts(m::MPSModel)
    if m.solver.provides_duals
	return m.reducedcosts
    else
	@warn("Reduced costs are not available from $(typeof(m.solver)). Try a different solver instead")
	return fill(NaN, m.ncol)
    end
end

function MPB.getconstrduals(m::MPSModel)
    if m.solver.provides_duals
	return 	m.duals
    else
	@warn("Constraint duals are not available from $(typeof(m.solver)). Try a different solver instead")
	return fill(NaN, m.nrow)
    end
end

function MPB.getsense(m::MPSModel)
    m.sense
end

function MPB.setsense!(m::MPSModel, sense::Symbol)
    m.sense = sense
end

function MPB.setvartype!(m::MPSModel, t::Vector{Symbol})
    m.colcat = t
end


text:
using LinearAlgebra
using Roots
using Intervals

"""
    GammaEqn

Save data for solving for gamma
"""
struct GammaEqn{T}
    ϵ::T
    λ::Array{T, 1}
    m::Int
end
GammaEqn(e::Int, l::AbstractArray, m::Int) = GammaEqn(Float64(e), l, m)
function GammaEqn(ϵ::Union{Real, Int}, λ)
    if !all( λ .>= 0 )
        throw(ArgumentError("Eigenvalues must be nonnegative"))
    else
        return GammaEqn(ϵ, λ, length(λ))
    end
end
function (eq::GammaEqn)(expr, γ)
    ϵ = eq.ϵ
    λ = eq.λ
    m = eq.m
    z = γ[1]
    expr[1] = (ϵ^2 .- 0.5*sum(λ)) .* z - m + 0.5 * sum(sqrt.(z.^2 .* λ.^2 .+ 4 .* λ .* z))
end

"""
    WassersteinShrinkage(target, eps)

Computes the Wasserstein-DR precision matrix of X wrt the Stein loss logdet(C*P) - <C*P> - p with the radius eps.
"""
struct WassersteinShrinkage <:CovarianceEstimator
    eps::Float64
end

function cov(sre::WassersteinShrinkage, X::AbstractMatrix{<:Real};
             dims::Int=1, mean=nothing)
    inv_cov = prec(sre, X)
    return inv(inv_cov)
end

function prec(sre::WassersteinShrinkage, X::AbstractMatrix{<:Real};
        dims::Int=1, mean=nothing)
    return wasserstein_pm(X, sre.eps)
end

"""
    linear_shrinkage(::DiagonalUnitVariance, Xc, S, λ, n, p, corrected)

Compute the shrinkage estimator where the target is a `DiagonalUnitVariance`.
"""
function wasserstein_pm(X, ϵ)
    C = Symmetric(cov(X))
    λ, v = eigen(C)
    m = length(λ)
    λ = max.(1e-10, λ)
    eq = GammaEqn(ϵ, λ)
    ϕ(z) = (ϵ^2 .- 0.5*sum(λ)) .* z - m + 0.5 * sum(sqrt.(z.^2 .* λ.^2 .+ 4 .* λ .* z))
    bounds = γ_bounds(eq)
    ϕ.(bounds)
    γ = find_zero(ϕ, bounds, Roots.Bisection())
    u = λ*γ
    # x = γ .* (1 .- 2. ./ (1 .+ sqrt.(1 .+ 4 ./ u)))
    x = γ.*(1 .- 0.5 *(sqrt.(max.(0., u.^2 .+ 4*u)) .- u))
    for i = 1:length(λ)
        # NOTE: Numerically stable alternative for large values
        if λ[i]*γ > 1e8
            x[i] = γ .* (1 .- 2. ./ (1 .+ sqrt.(1 .+ 4 ./ λ[i]*γ)))
            # x[i] = γ*(1 - 0.5 *(sqrt((γ*λ[i])^2 + 4*λ[i]*γ)) - λ[i]*γ)
        end
    end
    est = v*diagm(x)*v'
    return Symmetric(est)
end
function γ_bounds(eq::GammaEqn)
    m = eq.m
    L = eq.λ[end]
    ϵ = eq.ϵ
    # NOTE: Lower bound from original paper has an error?
    lower = 1e-15
    upper = min(m/ϵ^2, 1/(ϵ + 1e-15)*sqrt(sum(eq.λ.^(-1))))
    return lower, upper
end