text:
"""
    NCT

Data structure containing information needed to write a NetCDF file.

"""
struct NCT
    lon
    lat
    data
    meta
end

"""
    NCvar

Data structure containing information needed to write a variable to NetCDF files.

Instead of loading the data into memory, one can provide a list of input file names (see `Bindata`).
"""
struct NCvar
    name::String
    units::String
    dims
    values
    atts::Union{Dict,Nothing}
    backend::Module
end

function NCvar(name::String, units::String, dims::NCvar, values, atts::Union{Dict,Nothing}, backend::Module)
    return NCvar(name, units, [dims], values, atts, backend)
end

"""
    replacevalues(vals,ncvar::NCvar)

Helper function: Replaces the "values" attribute of an ncvar without changing
anything else. Used to apply land mask.
"""
replacevalues(vals,ncvar::NCvar) = NCvar(ncvar.name,ncvar.units,ncvar.dims,vals,ncvar.atts,ncvar.backend)

"""
    BinData

Data structure containing a string or an array of strings (NetCDF
    file names) as well as metadata needed to read a file.
"""
struct BinData # Pointer to data stored in binary files- contains info needed to read in
    fnames::Union{Array{String},String}
    precision::Type
    iosize::Tuple
    fldidx::Int
end

"""
    BinData(fnames::Union{Array{String},String},precision::Type,iosize::Tuple)

Construct a BinData struct for files that contain one field.
"""
function BinData(fnames::Union{Array{String},String},precision::Type,iosize::Tuple)
    return BinData(fnames,precision,iosize,1)
end

function BinData(fnames::Union{Array{String},String},precision::Type,iosize::Tuple,meta::Dict,fldname::String)
    return BinData(fnames,precision,iosize,findfirst(meta["fldList"] .== fldname)[2])
end

"""
    NCData

Data structure containing a string or an array of strings (file names) of
    NetCDF files as well as information needed to read a file.
"""
struct NCData
    fname::AbstractString
    varname::AbstractString
    backend::Module
    precision::Type
end

"""
    TileData{T}

Data structure containing either a `MeshArray` struct or `BinData` struct (see `vals`),
    `MeshArray` structs describing the tile layout (`tileinfo`), and other information for
    reading/writing tile data.
"""
struct TileData{T}
    vals::T
    tileinfo::Dict
    tilesize::Tuple
    precision::Type
    numtiles::Int
end

"""
    TileData(vals,tilesize::Tuple)

Construct a TileData struct. First generate the tileinfo, precision, and numtiles attributes.
"""
function TileData(vals,tilesize::Tuple,grid::gcmgrid)
    ni,nj = tilesize
    if isa(vals,BinData)
        prec = vals.precision
    else
        prec = eltype(vals)
    end
    τ=Tiles(grid,ni,nj); TileMap=MeshArray(grid,prec);
    [TileMap[τ[ii].face][τ[ii].i,τ[ii].j].=ii for ii in 1:length(τ)]
    tileinfo = Dict("tileNo" => TileMap)
    gridvars=GridLoad(grid)
    tileinfo["XC"]=gridvars.XC
    tileinfo["YC"]=gridvars.YC
    return TileData(vals,tileinfo,tilesize,prec,Int(maximum(tileinfo["tileNo"])))
end
TileData(vals,tilesize::Tuple,grid::String="LatLonCap") = TileData(vals,tilesize::Tuple,GridSpec(grid))
"""
    replacevalues(vals,td::TileData)

Helper function: Replaces the "values" attribute of a TilData struct without changing
anything else. Used to apply land mask.
"""
replacevalues(vals,td::TileData) = TileData(vals,td.tileinfo,td.tilesize,td.precision,td.numtiles)

import Base: getindex

"""
    getindex(v::NCvar,i::Integer)

Retrive an NCvar at time step i.

Currently only works when values is a BinData, NCdata, or array of time steps.
"""
function getindex(v::NCvar,i::Integer)
    NCvar(v.name,v.units,v.dims,v.values[i],v.atts,v.backend)
end

"""
    getindex(d::BinData,i::Integer)

 Return a new BinData object with the ith file, this being time index i.
"""
function getindex(d::BinData,i::Integer)
    if isa(d.fnames,Array)
        newfnames = [d.fnames[i]]
    elseif i == 1
        newfnames = d.fnames
    end

    BinData(newfnames,d.precision,d.iosize)
end

"""
    getindex(d::NCData,i)

Retrieve the ith time step from the NetCDF file.
"""
function getindex(d::NCData,i)
    readncdata(d,i)
end

"""
    addDim(ds::NCDatasets.Dataset,dimvar::NCvar) # NCDatasets

Add a dimension to a `NCDatasets.Dataset`
"""
function addDim(ds::NCDatasets.Dataset,dimvar::NCvar)
    defDim(ds,dimvar.name,dimvar.dims[1])
end

"""
    addDim(dimvar::NCvar)

Add a dimension to a NetCDF file using `NetCDF.jl`
"""
function addDim(dimvar::NCvar) #NetCDF
    dimvals = dimvar.values[:]
    NcDim(dimvar.name,collect(dimvals),
            atts = merge(Dict(("units" =>dimvar.units)),dimvar.atts),
            unlimited = dimvar.dims==Inf)
end

"""
    addVar(ds::NCDatasets.Dataset,field::NCvar)

Add a variable to a NetCDF file using `NCDatasets.jl`
"""
function addVar(ds::NCDatasets.Dataset,field::NCvar)
    if ~isempty(field.units)
        atts = merge(Dict(("units" =>field.units)),field.atts)
    else
        atts = field.atts
    end
    if isa(field.values,Array)
        if isa(field.values[1],Array)
            prec = typeof(field.values[1][1])
        else
            prec = typeof(field.values[1])
        end
        fieldvar = defVar(ds,field.name,prec,
                            tuple([f.name for f in field.dims]...),attrib=atts)
    else
        fieldvar = defVar(ds,field.name,field.values.precision,
                            tuple([f.name for f in field.dims]...),attrib=atts)
    end
end

"""
    addVar(field::NCvar,dimlist::Array{NetCDF.NcDim})

Add a variable with dimensions dimlist to a NetCDF file using `NetCDF.jl`
"""
function addVar(field::NCvar,dimlist::Array{NetCDF.NcDim})
    if ~isempty(field.units)
        atts = merge(Dict(("units" =>field.units)),field.atts)
    else
        atts = field.atts
    end
    if isa(field.values,Array)
        prec = eltype(field.values)
    else
        prec = field.values.precision
    end
    fieldvar = NcVar(field.name,dimlist,
                        atts = atts,
                        t = prec)
end

"""
    addVar(field::NCvar})

Add a variable and its dimensions to a NetCDF file using `NetCDF.jl`
"""
function addVar(field::NCvar)
    dimlist = addDim.(field.dims)
    if ~isempty(field.units)
        attributes = merge(Dict(("units" =>field.units)),field.atts)
    else
        attributes = field.atts
    end
    fieldvar = NcVar(field.name,dimlist,
        atts = merge(Dict(("units" =>field.units)),field.atts))
end

"""
    addData(v::Union{NCDatasets.CFVariable,NetCDF.NcVar},var::NCvar)

Fill variable with data in netcdf file using either `NCDatasets.jl`
or `NetCDF.jl`
"""
function addData(v::Union{NCDatasets.CFVariable,NetCDF.NcVar,Array},var::NCvar;startidx=1,land_mask=nothing)
    isBinData = isa(var.values,BinData)
    isNCData = isa(var.values,NCData)
    isTileData = isa(var.values,TileData)
    if isTileData
        isBinData = isa(var.values.vals,BinData)
    end
    nsteps = getnsteps(var.values)
    ndims = length(var.dims)
    if hastimedim(var)
        ndims = ndims-1
    end

    if isTileData
        if isnothing(land_mask)
            gridvars = GridLoad(var.values.tileinfo["XC"].grid)
            if any(occursin.("_s",getfield.(var.dims,:name)))
                land_mask = gridvars["hFacS"]
            elseif any(occursin.("_w",getfield.(var.dims,:name)))
                land_mask = gridvars["hFacW"]
            else
                land_mask = gridvars["hFacC"]
            end
            for f in land_mask.fIndex
                for d in 1:size(land_mask,2)
                    land_mask[f,d][land_mask[f,d].==0] .= NaN
                    land_mask[f,d][land_mask[f,d].>0] .= 1
                end
            end
        end
    end

    if (isa(var.values,Array) && isa(var.values[1],Number)) || ~hastimedim(var)
        if isa(var.values,Array)
            v0 = var.values
        else
            v0 = readdata(var.values,:)
        end

        if ~isnothing(land_mask)
            if length(size(v0)) == 1 && length(size(land_mask)) == 2
                v0 = v0 .* land_mask[:,1]
            elseif length(size(v0)) == 2 && length(size(land_mask)) == 3
                v0 = v0 .* land_mask[:,:,1]
            else
                v0 = v0 .* land_mask
            end
        end
        if isTileData
            tmpvar = replacevalues(replacevalues(v0,var.values),var)
            writetiles.(v,Ref(tmpvar),1:var.values.numtiles,Ref(:),Ref(land_mask))
        else
            ndims = length(var.dims)
            checkdims(v0,var::NCvar)
            if ndims == 1
                v[:] = v0
            elseif ndims == 2
                v[:,:] = v0
            elseif ndims == 3
                v[:,:,:] = v0
            elseif ndims == 4
                v[:,:,:,:] = v0
            end
        end

    elseif isBinData || isNCData || isTileData || isa(var.values[1],Array) # Binary files or array of timesteps

        if isBinData && ~ isTileData
            if isa(var.values.fnames,Array)
                fnames = var.values.fnames
            else
                fnames = [var.values.fnames]
            end
        end
            for i = startidx:nsteps
                if isBinData || isNCData || isTileData
                    v0 = readdata(var.values,i)
                else
                    v0 = var.values[i]
                end

                if ~isnothing(land_mask)
                    if length(size(v0)) == 1 && length(size(land_mask)) == 2
                        v0 = v0 .* land_mask[:,1]
                    elseif length(size(v0)) == 2 && length(size(land_mask)) == 3
                        v0 = v0 .* land_mask[:,:,1]
                    else
                        v0 = v0 .* land_mask
                    end
                end
                if isTileData
                    tmpvar = replacevalues(replacevalues(v0,var.values),var)
                    writetiles.(v,Ref(tmpvar),1:var.values.numtiles,Ref(i),Ref(land_mask))
                else
                    checkdims(v0,var)
                    if ndims == 0
                        v[i] = v0
                    elseif ndims == 1
                        v[:,i] = v0
                    elseif ndims == 2
                        v[:,:,i] = v0
                    elseif ndims == 3
                        v[:,:,:,i] = v0
                    end
                end
            end
    else
        print("Unrecognized values")
    end
    return nothing
end

"""
    writetiles(v,var,tilenum,timeidx=1)

Helper function for writing a tile to a NetCDF file.
"""
function writetiles(v,var,tilenum,timeidx=1,land_mask=nothing)
    if isa(v,Array)
        if var.backend == NCDatasets
            v = v[findfirst(isequal(var.name),name.(v))]
        else
            varnames = [vtmp.name for vtmp in v]
            v = v[findfirst(isequal(var.name),varnames)]
        end
    end
    tileinfo = var.values.tileinfo; tilesize = var.values.tilesize
    if isa(var.values.vals,MeshArray) || isa(var.values.vals,MeshArrays.gcmfaces)
        v0 = var.values.vals
    else
        v0 = readdata(var.values,timeidx)
    end

    v0 = gettile(v0,tileinfo,tilesize,tilenum)
    checkdims(v0,var::NCvar)
    if hastimedim(var)
        numdims = ndims(v0)
        if numdims == 1
            v[:,timeidx] = v0
        elseif numdims == 2
            v[:,:,timeidx] = v0
        elseif numdims == 3
            v[:,:,:,timeidx] = v0
        end
    else
        v[:] = v0
    end

end

"""
    addDimData(v::Union{NCDatasets.CFVariable,NetCDF.NcVar,Array},var::NCvar)

Add dimension data to predefined dimensions in a NetCDF file.
"""
function addDimData(ds,dimvar::NCvar)
    atts = merge(Dict(("units" =>dimvar.units)),dimvar.atts)
    defVar(ds,dimvar.name,dimvar.values[:],(dimvar.name,),attrib=atts)
end

"""
    createfile(filename, field::Union{NCvar,Dict{String,NCvar}}, README;
                fillval=NaN, missval=NaN, itile=1, ntile=1)

Create NetCDF file and add variable + dimension definitions
using either `NCDatasets.jl` or `NetCDF.jl`
"""
function createfile(filename, field::Union{NCvar,Dict}, rdm="";
                    fillval=NaN, missval=NaN, itile=1, ntile=1, attribs=nothing)

    if isa(field,Dict)

        dims = getdims(field::Dict)
        field = collect(values(field))

        fldnames = getfield.(field[findall(hastimedim.(field))],:name)
        backend = field[1].backend
    else
        dims = field.dims
        fldnames = field.name
        backend = field.backend
        if backend == NCDatasets
            field = [field]
        end
    end

    #test_climatology_bounds=sum(occursin.(fldnames,"climatology_bounds"))

    if isa(fldnames,Array)
        fldnamestring = join(filter(x -> x !== "climatology_bounds", fldnames),",")
    else
        fldnamestring = fldnames
    end

    README=deepcopy(rdm)
    if ~isempty(README) && ~occursin(fldnamestring*" -- ",README[1])
        README[1] = fldnamestring*" -- "*README[1]
    end

    if ~isnothing(attribs)
        fillval = pop!(attribs,"_FillValue",NaN)
        missingval = pop!(attribs,"missing_value",NaN)
        haskey(attribs,"itile") ? itile = pop!(attribs,"itile") : nothing
        haskey(attribs,"ntile") ? ntile = pop!(attribs,"ntile") : nothing
        description = [pop!(attribs,k,"") for k in ["description","A","B","C","D","E","F","G","H","I","J"]]
        description = description[.~isempty.(description)]
        if isempty(README)
            README = description
        elseif !(description == README)
            README = vcat(README,description)
        end
        pop!(attribs,"date",nothing)
        if isempty(attribs); attribs=nothing; end
    end

    if ~isempty(README)
        if isa(README,Array) && length(README) > 1
            description = vcat("description" => README[1],
            [string(Char(65+(i-2))) => README[i] for i in 2:length(README)])
        else
            description = ["description" => isa(README,Array) ? README[1] : README]
        end
    else
        description = ""
    end

    file_atts = vcat(["date" => Dates.format(today(),"dd-u-yyyy"),
    "Conventions" => "CF-1.6"], "NCTiles" => NCTiles.version())

    if ~isempty(description)
        file_atts = vcat(file_atts,description)
    end
    file_atts = vcat(file_atts,
    [#string(Char(65+length(README)-1)) => "file created using NCTiles.jl",
    "_FillValue" => fillval,
    "missing_value" => missval,
    "itile" => itile,
    "ntile" => ntile])

    if ~isnothing(attribs)
        file_atts = vcat(file_atts,[k => attribs[k] for k in keys(attribs)])
    end

    if backend == NCDatasets
        ds =  Dataset(filename,"c",attrib=file_atts)
        dimlist = addDim.(Ref(ds),dims)
        fieldvar = addVar.(Ref(ds), field)
        if length(fieldvar) == 1
            fieldvar = fieldvar[1]
        end
        return ds,fieldvar,dimlist
    elseif backend == NetCDF
        dimlist = addDim.(dims)
        #test_climatology_bounds ? addVar(f,dimlist[dimidx]) : nothing

        if isa(field,Array)
            dimnames = getfield.(dims,:name)
            field = filter(f ->~any(f.name.==dimnames),field) # need to get dimlist for each variable
            nflds = length(field)
            fieldvar = Array{NcVar}(undef,nflds)
            for i in 1:nflds
                f = field[i]
                dimidx = [findfirst(d -> d .== dim.name,dimnames) for dim in f.dims]
                fieldvar[i] = addVar(f,dimlist[dimidx])
            end
        else
            fieldvar = addVar(field,dimlist)
        end
        file_atts = Dict(file_atts)
        return NetCDF.create(filename,fieldvar, gatts = file_atts),fieldvar,dimlist
    end

end

import Base: write
using Pkg

write(nct::NCT) = write_2D_clim(nct.lon,nct.lat,nct.data,nct.meta)

function write_2D_clim(lon,lat,data,meta)
	nc=NCTiles.NetCDF #backend

	#1. pre-define dimensions
	lon_c=NCvar("lon","degrees_east",size(lon,1),lon[:,1],Dict("long_name" => "longitude"),nc)
	lat_c=NCvar("lat","degrees_north",size(lat,2),lat[1,:],Dict("long_name" => "latitude"),nc)

	yrng=meta.yearrange
	tim_units="days since $(yrng[1])-1-1 0:0:0"
	tim_values=NCTiles.DateTime.(yrng[1],collect(1:12),16)
	tim_bounds=permutedims([DateTime.(yrng[1],collect(1:12)) [DateTime.(yrng[2],collect(2:12));DateTime.(yrng[2]+1,1)]],(2,1))

    nt=length(tim_values)
    t_values=Float64.(collect(1:nt))
	t=NCvar("t", "1", nt, t_values, Dict("units" => "1","long_name" => "Time coordinate"), nc)
	dims = [lon_c, lat_c, t]

    globalattribs=Dict("_FillValue" => NaN,"missing_value" => NaN,"itile" => 1.0,"ntile" => 1.0)

	#2. put climatology from Array into NCvar with meta-data
	ncv = NCvar(meta.name,meta.units,dims,data, Dict("long_name" => meta.long_name),nc)

	#3. prepare inputs to `write` method
	msk=Float64.(isnan.(data[:,:,1]))
	land=NCvar("land","1.0",dims[1:2],msk,Dict(["long_name" => "land mask",
				"standard_name" => "land_binary_mask"]), nc)
	#msk=Float64.(isnan.(data))
	#land=NCvar("land","1.0",dims,data,Dict("long_name" => "land mask",
	#			"standard_name" => "land_binary_mask"]), nc)

	writevars = Dict([meta.name => ncv,"land" => land,
			"lon" => lon_c, "lat" => lat_c, "t" => dims[3]])
#                "dep" => dep,"thic" => thic,"area" => area,

	#4. create and instantiate the new file
	rm(meta.outputfile,force=true)
	write(writevars,meta.outputfile;README=meta.README,globalattribs=globalattribs)

	#5. add climatology_bounds and tim afterwards
	ds = Dataset(meta.outputfile,"a")
	defDim(ds,"tcb",2)

	v = defVar(ds,"climatology_bounds",Float32,("tcb","t"),
	attrib = Dict("units" => tim_units, "long_name" => "climatology_bounds"))
	v[:,:] = tim_bounds

	v = defVar(ds,"tim",Float32,("t",), attrib = Dict("units" => tim_units,
    "climatology" => "climatology_bounds", "long_name" => "time", "standard_name" => "time"))
	v[:] = tim_values

	close(ds)
end

"""
    write(myfld::NCvar,savename::String;README="",globalattribs=Dict())

Creates NetCDF file and writes myfld and all its dimensions to the file.
"""
function write(myfld::NCvar,savename::String;README="",globalattribs=Dict())
    if hastiledata(myfld) # Create one file for each tile

        numtiles = myfld.values.numtiles
        savenames = savename*".".*lpad.(string.(1:numtiles),4,"0").*".nc"

        datasets = [createfile(savenames[tidx],myfld,README, itile = tidx, ntile = length(savenames), attribs = globalattribs) for tidx in 1:length(savenames)]

        ds = [x[1] for x in datasets]
        fldvars = [x[2] for x in datasets]

        addData(fldvars,myfld)

        if myfld.backend == NCDatasets
            dims = myfld.dims

            for dim in dims
                addDimData.(ds,Ref(dim))
            end

            close.(ds)
        else
            if pkg_version("NetCDF") < v"0.10"
                NetCDF.close.(ds)
            else
                finalize.(ds)
            end
        end
    else # Create one file
        ds,fldvar,dimlist = createfile(savename,myfld,README,attribs=globalattribs)

        # Add field and dimension data
        addData(fldvar,myfld)

        if myfld.backend == NCDatasets
            addDimData.(Ref(ds),myfld.dims)

            # Close the file
            close(ds)
        else
            if pkg_version("NetCDF") < v"0.10"
                NetCDF.close(ds)
            else
                finalize(ds)
            end
        end
    end
end

"""
    write(myflds::Dict,savename::String;README="",globalattribs=Dict())

Creates NetCDF file and writes myflds and all their dimensions to the file.
"""
function write(myflds::Dict{String,NCvar},savename::String;README="",globalattribs=Dict())

    if hastiledata(myflds)

        fldnames = collect(keys(myflds))
        tilefld = myflds[fldnames[findfirst([hastiledata(myflds[f]) for f in fldnames])]]
        numtiles = tilefld.values.numtiles

        landidx = findfirst(get.([myflds[f].atts for f in fldnames],"standard_name","none").=="land_binary_mask")
        if ~isnothing(landidx); land_mask = myflds[fldnames[landidx]].values; else land_mask = nothing; end
        if isa(land_mask,TileData); land_mask = land_mask.vals; end
        savenames = savename*".".*lpad.(string.(1:numtiles),4,"0").*".nc"

        datasets = [createfile(savenames[tidx],myflds,README, itile = tidx, ntile = length(savenames), attribs = globalattribs) for tidx in 1:length(savenames)]

        ds = [x[1] for x in datasets]
        fldvars = [x[2] for x in datasets]

        for k in keys(myflds)
            if isa(myflds[k].values,TileData)
                addData(fldvars,myflds[k],land_mask = land_mask)
            else
                if myflds[k].backend == NCDatasets
                    tmpfldvars = [fv[findfirst(isequal(k),name.(fv))] for fv in fldvars]
                else
                    tmpfldvars = [fv[findfirst(isequal(k),getproperty.(fv,:name))] for fv in fldvars]
                end

                addData.(tmpfldvars,Ref(myflds[k]))
            end
        end
        if isa(ds[1],NCDatasets.Dataset)
            dims = unique(vcat([myflds[v].dims for v in keys(myflds)]...))
            dims = filter( d -> isa(d,NCvar),dims)

            for dim in dims
                addDimData.(ds,Ref(dim))
            end

            close.(ds)
        else
            if pkg_version("NetCDF") < v"0.10"
                NetCDF.close.(ds)
            else
                finalize.(ds)
            end
        end

    else

        dims = getdims(myflds)

        ## Here down to move to other branch- high level API for writing
        ds,fldvar,dimlist = createfile(savename,myflds,README,attribs=globalattribs)

        # Add field data
        for k in keys(myflds)
            addData(ds[k],myflds[k])
        end

        if isa(ds,NCDatasets.Dataset)
            # Only insert data for dims with data
            dims = filter( d -> ~isempty(d.values[:]),dims)

            # Add dimension data
            addDimData.(Ref(ds),dims)

            # Close the file
            close(ds)
        else
            if pkg_version("NetCDF") < v"0.10"
                NetCDF.close(ds)
            else
                finalize(ds)
            end
        end
    end
end

write(myflds::Dict,savename::String;README="",globalattribs=Dict()) = write(Dict{String,NCvar}(myflds),savename;README="",globalattribs=Dict())


text:
const XSIMD_DICT = Dict(
    (f => (:xsimdwrap, f) for f ∈ float_to_float)...,
    (f => (:xsimdwrap, f) for f ∈ floatfloat_to_float)...,
    (f => (:xsimdwrap, f) for f ∈ float_to_floatfloat)...,
    :^ => (:xsimdwrap, :pow)
)

"""
Arguments are
@xvectorize Type UnrollFactor forloop
The default type is Float64, and default UnrollFactor is 1 (no unrolling).
"""

for gcpreserve ∈ (true,false)
    if gcpreserve
        macroname = :xvectorize
    else
        macroname = Symbol(:xvectorize, :_unsafe)
    end
    @eval macro $macroname(expr)
        if @capture(expr, for n_ ∈ 1:N_ body__ end)
            q = vectorize_body(N, Float64, 1, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(A_) body__ end)
            q = vectorize_body(:(length($A)), Float64, 1, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(args__) body__ end)
            q = vectorize_body(:(min($([:(length($a)) for a ∈ args]...))), Float64, 1, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        else
            throw("Could not match loop expression.")
        end
        esc(q)
    end
    @eval macro $macroname(type, expr)
        if @capture(expr, for n_ ∈ 1:N_ body__ end)
            q = vectorize_body(N, type, 1, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(A_) body__ end)
            q = vectorize_body(:(length($A)), type, 1, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(args__) body__ end)
            q = vectorize_body(:(min($([:(length($a)) for a ∈ args]...))), type, 1, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        else
            throw("Could not match loop expression.")
        end
        esc(q)
    end
    @eval macro $macroname(unroll_factor::Integer, expr)
        if @capture(expr, for n_ ∈ 1:N_ body__ end)
            q = vectorize_body(N, Float64, unroll_factor, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(A_) body__ end)
            q = vectorize_body(:(length($A)), Float64, unroll_factor, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(args__) body__ end)
            q = vectorize_body(:(min($([:(length($a)) for a ∈ args]...))), Float64, unroll_factor, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        else
            throw("Could not match loop expression.")
        end
        esc(q)
    end
    @eval macro $macroname(type, unroll_factor::Integer, expr)
        if @capture(expr, for n_ ∈ 1:N_ body__ end)
            q = vectorize_body(N, type, unroll_factor, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(A_) body__ end)
            q = vectorize_body(:(length($A)), type, unroll_factor, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(args__) body__ end)
            q = vectorize_body(:(min($([:(length($a)) for a ∈ args]...))), type, unroll_factor, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, :xsimdwrap)
        else
            throw("Could not match loop expression.")
        end
        esc(q)
    end
    @eval macro $macroname(type, mod::Union{Symbol,Module}, expr)
        if @capture(expr, for n_ ∈ 1:N_ body__ end)
            q = vectorize_body(N, type, 1, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, mod, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(A_) body__ end)
            q = vectorize_body(:(length($A)), type, 1, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, mod, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(args__) body__ end)
            q = vectorize_body(:(min($([:(length($a)) for a ∈ args]...))), type, 1, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, mod, :xsimdwrap)
        else
            throw("Could not match loop expression.")
        end
        esc(q)
    end
    @eval macro $macroname(type, mod::Union{Symbol,Module}, unroll_factor::Integer, expr)
        if @capture(expr, for n_ ∈ 1:N_ body__ end)
            q = vectorize_body(N, type, unroll_factor, n, body, XSIMD_DICT, SIMDPirates.Vec, mod, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(A_) body__ end)
            q = vectorize_body(:(length($A)), type, unroll_factor, n, body, XSIMD_DICT, SIMDPirates.Vec, mod, :xsimdwrap)
        elseif @capture(expr, for n_ ∈ eachindex(args__) body__ end)
            q = vectorize_body(:(min($([:(length($a)) for a ∈ args]...))), type, unroll_factor, n, body, XSIMD_DICT, SIMDPirates.Vec, $gcpreserve, mod, :xsimdwrap)
        else
            throw("Could not match loop expression.")
        end
        esc(q)
    end
end

ext: jl


text:
using GridWorlds
using Test
using Random
using ReinforcementLearningBase

ENVS = [EmptyGridWorld, FourRooms, GoToDoor, DoorKey, CollectGems, DynamicObstacles]
ENVS_RLBASE = [EmptyGridWorld, FourRooms, GoToDoor, DoorKey, CollectGems, DynamicObstacles]
ACTIONS = [TURN_LEFT, TURN_RIGHT, MOVE_FORWARD]

@testset "GridWorlds.jl" begin
    for Env in ENVS
        @testset "$(Env)" begin
            env = Env()
            @test typeof(env.agent_pos) == CartesianIndex{2}
            @test typeof(env.agent.dir) <: LRUD
            @test size(env.world.grid, 1) == length(env.world.objects)
            @test 1 ≤ env.agent_pos[1] ≤ size(env.world.grid, 2)
            @test 1 ≤ env.agent_pos[2] ≤ size(env.world.grid, 3)

            for _=1:1000
                env = env(rand(ACTIONS))
                @test 1 ≤ env.agent_pos[1] ≤ size(env.world.grid, 2)
                @test 1 ≤ env.agent_pos[2] ≤ size(env.world.grid, 3)
                @test env.world[WALL, env.agent_pos] == false
                view = get_agent_view(env)
                @test typeof(view) <: BitArray{3}
                @test size(view,1) == length(env.world.objects)
                @test size(view,2) == 7
                @test size(view,3) == 7
            end
        end
    end

    for Env in ENVS_RLBASE
        @testset "$(Env) RLBase API" begin
            env = Env()

            @test get_reward(env) == 0.0
            @test get_terminal(env) == false
            if Env == GoToDoor
                @test get_state(env) == (get_agent_view(env), env.target)
            else
                @test get_state(env) == get_agent_view(env)
            end

            π = RandomPolicy(env)
            total_reward = 0.0

            while !get_terminal(env)
                action = π(env)
                env(action)
                total_reward += get_reward(env)
            end

            if Env == CollectGems
                @test total_reward == env.num_gem_init * env.gem_reward
            elseif Env == EmptyGridWorld || Env == FourRooms || Env == DoorKey
                @test total_reward == env.goal_reward
            elseif Env == GoToDoor
                @test (total_reward == env.target_reward || total_reward == env.penalty)
            elseif Env == DynamicObstacles
                @test (total_reward == env.obstacle_reward || total_reward == env.goal_reward)
            end
        end
    end
end


text:
using ImageCore
using FileIO

function image_from_file(filename::String, width, height)
    raw_image = reinterpret(UInt8, load(filename))
    vx_input_image = VX.CreateImage(ctx, width, height, VX.DF_IMAGE_RGB)

    image_region = Ref{VX.vx_rectangle_t}()
    image_region[] = VX.vx_rectangle_t(0, 0, width, height)
    image_layout = Ref{VX.vx_imagepatch_addressing_t}()
    image_layout[] = VX.vx_imagepatch_addressing_t(stride_x = 3, stride_y = width * 3)
    VX.CopyImagePatch(vx_input_image, image_region, 0, image_layout, raw_image,
                      VX.WRITE_ONLY, VX.MEMORY_TYPE_HOST)
    return vx_input_image
end

function file_from_image(image, filename::String, width, height)

    local_image = zeros(UInt8, 3, width, height)
    image_region = Ref{VX.vx_rectangle_t}()
    VX.GetValidRegionImage(image, image_region)

    image_layout = Ref{VX.vx_imagepatch_addressing_t}()
    map_id = Ref{VX.vx_map_id}()
    ptr = Ref{Ptr{UInt8}}()

    VX.MapImagePatch(image, image_region[], 0, map_id, image_layout, ptr,
                     VX.READ_ONLY, VX.MEMORY_TYPE_HOST, VX.NOGAP_X)

    dims = (3, Int(image_layout[].stride_y ÷ 3), height)
    mapped_vx_image = unsafe_wrap(Array, ptr[], dims)

    #copy our loaded image to VX mapped memory
    local_image .= mapped_vx_image

    if VX.UnmapImagePatch(image, map_id[]) == 0
        mapped_vx_image = nothing
    end

    normed = reinterpret.(N0f8, local_image)
    color_view = colorview(RGB, normed)

    save(filename, color_view)

    return color_view
end

function file_from_gray_image(image, filename::String, width, height)

    image_region = Ref{VX.vx_rectangle_t}()
    VX.GetValidRegionImage(image, image_region)

    image_layout = Ref{VX.vx_imagepatch_addressing_t}()
    map_id = Ref{VX.vx_map_id}()
    ptr = Ref{Ptr{UInt8}}()

    VX.MapImagePatch(image, image_region[], 0, map_id, image_layout, ptr,
                     VX.READ_ONLY, VX.MEMORY_TYPE_HOST, VX.NOGAP_X)

    dims = (Int(image_layout[].stride_y), height)
    mapped_vx_image = unsafe_wrap(Array, ptr[], dims)

    #copy from VX mapped memory
    local_image = mapped_vx_image[1:Int(image_layout[].dim_x), 1:Int(image_layout[].dim_y)]

    normed = reinterpret.(N0f8, local_image)
    color_view = colorview(Gray, normed)

    save(filename, color_view)

    if VX.UnmapImagePatch(image, map_id[]) == 0
        mapped_vx_image = nothing
    end

    return color_view
end


text:
"""
Solver data structure definition
"""
mutable struct Solver

    inst::Instance # instance to solve
    master::Model # master problem model
    sub::Model # subproblem model
    sol::Solution # corresponding solution

    """
    Solver object builder
    """
    function Solver(inst::Instance)

        this = new()
        this.inst = inst
        this.sol =  Solution(inst)
        return this

    end

end

"""
Generates and initializes the optimization model corresponding to the master problem
"""
function generate_master_model!(sv::Solver; relax_integrality::Bool = false)

    sv.master = new_model()

    # advanced constraint generation
    if args["method"] == "adv-constr-gen"

        # declaration of variables
        if relax_integrality
            @variable(sv.master, 0 <= x[a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)])] <= 1)
            @variable(sv.master, 0 <= t[1:sv.inst.n] <= 1)
        else
            @variable(sv.master, x[a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)])], Bin)
            @variable(sv.master, t[1:sv.inst.n], Bin)
        end

        # declaration of the objective function
        @expression(sv.master, distance, sum(sv.inst.D[a]*x[a] for a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)])))
        @objective(sv.master, Min, distance)

        # declaration of constraints
        @constraint(sv.master,
            force_artificial_arc,
            x[(sv.inst.F, sv.inst.I)] == 1)
        @constraint(sv.master,
            minimum_visited_airdromes,
            sum(x[a] for a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)])) >= sv.inst.k)
        @constraint(sv.master,
            unique_visit[i in 1:sv.inst.n],
            sum(x[a] for a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)]) if a[1] == i)
            + sum(x[a] for a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)]) if a[2] == i) == 2*t[i])
        @constraint(sv.master,
            touch_and_go[i in 1:sv.inst.n],
            sum(x[a] for a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)]) if a[1] == i)
            - sum(x[a] for a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)]) if a[2] == i) == 0)
        @constraint(sv.master,
            geographical_diversity[r in 1:sv.inst.m],
            sum(x[a] for a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)]) if sv.inst.R[a[1]] == r || sv.inst.R[a[2]] == r) >= 1)

    else
        # declaration of variables
        if relax_integrality
            @variable(sv.master, 0 <= x[a in sv.inst.feasible_arcs] <= 1)
        else
            @variable(sv.master, x[a in sv.inst.feasible_arcs], Bin)
        end

        if args["method"] == "sequential"
            @variable(sv.master, u[1:sv.inst.n] >= 0)

        elseif args["method"] == "single-flow"
            @variable(sv.master, q[a in sv.inst.feasible_arcs] >= 0)
            if relax_integrality
                @variable(sv.master, 0 <= φ[1:sv.inst.n] <= 1)
            else
                @variable(sv.master, φ[1:sv.inst.n], Bin)
            end
        end

        # declaration of the objective function
        @expression(sv.master, distance, sum(sv.inst.D[a]*x[a] for a in sv.inst.feasible_arcs))
        @objective(sv.master, Min, distance)

        # declaration of constraints
        @constraint(sv.master,
            minimum_visited_airdromes,
            sum(x[a] for a in sv.inst.feasible_arcs) >= sv.inst.k-1)
        @constraint(sv.master,
            unique_visit[i in 1:sv.inst.n],
            sum(x[a] for a in sv.inst.feasible_arcs if a[1] == i) + sum(x[a] for a in sv.inst.feasible_arcs if a[2] == i) <= 2)
        @constraint(sv.master,
            touch_and_go[i in 1:sv.inst.n],
            sum(x[a] for a in sv.inst.feasible_arcs if a[1] == i)
            - sum(x[a] for a in sv.inst.feasible_arcs if a[2] == i) == Int(i == sv.inst.I) - Int(i == sv.inst.F))
        @constraint(sv.master,
            geographical_diversity[r in 1:sv.inst.m],
            sum(x[a] for a in sv.inst.feasible_arcs if sv.inst.R[a[1]] == r || sv.inst.R[a[2]] == r) >= 1)

        if args["method"] == "sequential"
            @constraint(sv.master,
                increasing_order_of_visit[a in sv.inst.feasible_arcs],
                u[a[2]] - u[a[1]] + sv.inst.n*(1-x[a]) >= 1)

        elseif args["method"] == "single-flow"
            @constraint(sv.master,
                flow_limitation[a in sv.inst.feasible_arcs],
                (sv.inst.n-1)*x[a] - q[a] >= 0)
            @constraint(sv.master,
                total_flow,
                sum(q[a] for a in sv.inst.feasible_arcs if a[1] == sv.inst.I) - sum(φ[i] for i in setdiff(1:sv.inst.n, sv.inst.I)) == 0)
            @constraint(sv.master,
                flow_conservation[i in setdiff(1:sv.inst.n, sv.inst.I)],
                sum(q[a] for a in sv.inst.feasible_arcs if a[2] == i) - sum(q[a] for a in sv.inst.feasible_arcs if a[1] == i) - φ[i] == 0)
            @constraint(sv.master,
                virtual_demand[i in setdiff(1:sv.inst.n, sv.inst.I)],
                sum(x[a] for a in sv.inst.feasible_arcs if a[2] == i) - φ[i] == 0)
        end
    end

end

"""
Generates and initializes the optimization model corresponding to the subproblem
"""
function generate_sub_model!(sv::Solver; relax_integrality::Bool = false)

    sv.sub = new_model()

    # declaration of variables
    if relax_integrality
        @variable(sv.sub, 0 <= z[1:sv.inst.n] <= 1)
    else
        @variable(sv.sub, z[1:sv.inst.n], Bin)
    end

    @constraint(sv.sub, non_empty_subset, sum(z[i] for i in 1:sv.inst.n) >= 1)

    # advanced constraint generation
    if args["method"] == "adv-constr-gen"
        @variable(sv.sub, y[a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)])] >= 0)

        # declaration of constraints
        @constraint(sv.sub,
            linearization1[a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)])],
            y[a] - z[a[1]] <= 0)
        @constraint(sv.sub,
            linearization2[a in union(sv.inst.feasible_arcs, [(sv.inst.F, sv.inst.I)])],
            y[a] - z[a[2]] <= 0)

        # dummy constraint that will be updated
        sv.sub[:pivot] = @constraint(sv.sub, 0 <= 1)

    # standard constraint generation
    else
        @variable(sv.sub, y[a in sv.inst.feasible_arcs] >= 0)

        # declaration of constraints
        @constraint(sv.sub,
            linearization1[a in sv.inst.feasible_arcs],
            y[a] - z[a[1]] <= 0)
        @constraint(sv.sub,
            linearization2[a in sv.inst.feasible_arcs],
            y[a] - z[a[2]] <= 0)
        @constraint(sv.sub,
            linearization3[a in sv.inst.feasible_arcs],
            z[a[1]] + z[a[2]] - y[a] <= 1)
    end

end

"""
Updates the subproblem model by adding a subtour elimination constraint
"""
function add_subtour_constraint!(sv::Solver, z_cur::Dict{Int, Int}, nb_added_subtours::Int, p::Union{Int, Nothing} = nothing)

    # standard constraint generation
    if args["method"] == "std-constr-gen"
        sv.master[Symbol("subtour$(nb_added_subtours)")] = @constraint(sv.master,
            sum(sv.master[:x][a] for a in sv.inst.feasible_arcs if z_cur[a[1]] == 1 && z_cur[a[2]] == 1) - sum(z_cur[i] for i in 1:sv.inst.n) + 1 <= 0)

    # advanced constraint generation
    else
        sv.master[Symbol("subtour$(nb_added_subtours)")] = @constraint(sv.master,
            sum(sv.master[:x][a] for a in sv.inst.feasible_arcs if z_cur[a[1]] == 1 && z_cur[a[2]] == 1 && a[1] != sv.inst.I && a[2] != sv.inst.I)
            - sum(sv.master[:t][i] for i in setdiff(1:sv.inst.n, Set([p, sv.inst.I])) if z_cur[i] == 1) <= 0)
    end
end

"""
Updates the subproblem model's objective using the current solution of the master problem
"""
function update_sub_model_objective!(sv::Solver, x_cur::Dict{Tuple{Int, Int}, Float64}, t_cur::Union{Dict{Int, Float64}, Nothing} = nothing, p::Union{Int, Nothing} = nothing)

    # standard constraint generation
    if args["method"] == "std-constr-gen"
        sv.sub[:violation] = @expression(sv.sub, sum(x_cur[a]*sv.sub[:y][a] for a in sv.inst.feasible_arcs) - sum(sv.sub[:z][i] for i in 1:sv.inst.n) + 1)

    # advanced constraint generation
    else
        sv.sub[:violation] = @expression(sv.sub,
            sum(x_cur[a]*sv.sub[:y][a] for a in sv.inst.feasible_arcs if a[1] != sv.inst.I && a[2] != sv.inst.I)
            - sum(t_cur[i]*sv.sub[:z][i] for i in 1:sv.inst.n if i != p && i != sv.inst.I))
    end

    @objective(sv.sub, Max, sv.sub[:violation])

end

"""
(Advanced constraint generation ONLY) Updates the pivot p in the subproblem model
"""
function update_sub_model_pivot_constraint!(sv::Solver, p::Int)

    delete(sv.sub, sv.sub[:pivot])
    sv.sub[:pivot] = @constraint(sv.sub, sv.sub[:z][p] == 1)

end

"""
Runs the optimization
"""
function solve!(sv::Solver)

    # non-iterative methods => single MILP to solve
    if args["method"] == "sequential" || args["method"] == "single-flow"
        generate_master_model!(sv)
        if args["verbosity"] == 0
            stdoff()
        end
        optimize!(sv.master)
        if args["verbosity"] == 0
            stdon()
        end
        ln1("")

    # standard constraint generation
    elseif args["method"] == "std-constr-gen"
        nb_added_subtours = 0
        x_cur = Dict{Tuple{Int, Int}, Float64}()
        z_cur = Dict{Int, Int}()
        generate_master_model!(sv)
        generate_sub_model!(sv)

        while true
            # solve the master problem
            stdoff()
            optimize!(sv.master)
            stdon()
            ln1("Iteration n°$(nb_added_subtours+1) ($(nb_added_subtours) subtour constraints generated so far)")
            ln1("Master objective value: $(value(sv.master[:distance]))")

            # retrieve the solution and update the subproblem model
            for a in sv.inst.feasible_arcs
                x_cur[a] = value(sv.master[:x][a])
            end
            update_sub_model_objective!(sv, x_cur)

            # solve the subproblem
            stdoff()
            optimize!(sv.sub)
            stdon()
            ln2("Subproblem objective value: $(value(sv.sub[:violation]))")
            ln1("")

            # while loop stops where there is no more subtour
            if value(sv.sub[:violation]) <= 1e-8
                break
            end

            # if there is a subtour, update the master problem model
            nb_added_subtours += 1
            for i in 1:sv.inst.n
                z_cur[i] = round(Int, value(sv.sub[:z][i]))
            end
            add_subtour_constraint!(sv, z_cur, nb_added_subtours)
        end

    # advanced constraint generation
    else
        nb_iterations = 0
        nb_added_subtours = 0
        added_subtours_history = Tuple{Int, Dict{Int, Int}}[]
        x_cur = Dict{Tuple{Int, Int}, Float64}()
        t_cur = Dict{Int, Float64}()
        z_cur = Dict{Int, Int}()
        generate_sub_model!(sv, relax_integrality = true)

        # two steps: 1. LP master 2. MILP master
        for master_relax_state in [true, false]
            generate_master_model!(sv, relax_integrality = master_relax_state)

            if !master_relax_state
                ln1("> Phase 2: MILP master problem")
                # when starting Step 2, add all the constraints added during Step 1
                for st in 1:nb_added_subtours
                    add_subtour_constraint!(sv, added_subtours_history[st][2], st, added_subtours_history[st][1])
                end
            else
                ln1("> Phase 1: LP master problem")
            end

            while true
                nb_iterations += 1

                # solve the master problem
                stdoff()
                optimize!(sv.master)
                stdon()
                ln1("Iteration n°$(nb_iterations) ($(nb_added_subtours) subtour constraints generated so far)")
                ln1("Master objective value: $(value(sv.master[:distance]))")

                # retrieve the solution and update the subproblem model
                for a in sv.inst.feasible_arcs
                    x_cur[a] = value(sv.master[:x][a])
                end
                for i in 1:sv.inst.n
                    t_cur[i] = value(sv.master[:t][i])
                end

                # for each possible pivot, solve a subproblem
                found_new_subtour = false
                for p in setdiff(1:sv.inst.n, sv.inst.I)
                    update_sub_model_pivot_constraint!(sv, p)
                    update_sub_model_objective!(sv, x_cur, t_cur, p)

                    # solve the subproblem
                    stdoff()
                    optimize!(sv.sub)
                    stdon()
                    ln2("Subproblem with p = $p objective value: $(value(sv.sub[:violation]))")
                    if value(sv.sub[:violation]) >= 1e-8
                        nb_added_subtours += 1
                        for i in 1:sv.inst.n
                            z_cur[i] = round(Int, value(sv.sub[:z][i]))
                        end
                        # keep track of subtours detected during Step 2 (needed to initialize Step 2)
                        if master_relax_state
                            push!(added_subtours_history, (copy(p), copy(z_cur)))
                        end
                        add_subtour_constraint!(sv, z_cur, nb_added_subtours, p)
                        found_new_subtour = true
                    end
                end
                ln1("")

                # while loop stops where there is no more subtour
                if !found_new_subtour
                    break
                end
            end
        end
    end

    # save the solution
    save_solution!(sv)
end

"""
Saves the solution as a Solution object
"""
function save_solution!(sv)

    # in advanced constraint generation, the arc (F,I) is fictive and must be ignored
    if args["method"] == "adv-constr-gen"
        sv.sol.score = value(sv.master[:distance]) - sv.inst.D[(sv.inst.F, sv.inst.I)]
    else
        sv.sol.score = value(sv.master[:distance])
    end

    # save selection rate for each feasible arc
    for a in sv.inst.feasible_arcs
        sv.sol.arc_selection[a] = min(max(0, value(sv.master[:x][a])), 1)
    end

    # in advanced constraint generation, the arc (F,I) is fictive and must be ignored
    if (sv.inst.F, sv.inst.I) in sv.inst.feasible_arcs
        sv.sol.arc_selection[(sv.inst.F, sv.inst.I)] = 0
    end

end

max_stars_repo_name: antcrepin/flight-route-optimizer

max_issues_repo_name: antcrepin/flight-route-optimizer

max_forks_repo_name: antcrepin/flight-route-optimizer


text:
abstract type AbstractFisheryDomain <: Any end
abstract type DiscreteFisheryDomain <: AbstractFisheryDomain end

"""
    GriddedFisheryDomain{T<:Real, Ti<:Integer} <: DiscreteFisheryDomain
        origin::Tuple{T, T}
        antipode::Tuple{T, T}
        n::Tuple{Ti, Ti}
        locs::Array{Tuple{T, T}, 2}
        distances::Array{T, 2}

Type to describe a rectangular, gridded domain.
"""
struct GriddedFisheryDomain{T<:Real, Ti<:Integer} <: DiscreteFisheryDomain
    origin::Tuple{T, T}
    antipode::Tuple{T, T}
    n::Tuple{Ti, Ti}
    locs::Array{Tuple{T, T}, 2}
    distances::Array{T, 2}
    steps::Tuple{T, T}

    function GriddedFisheryDomain{T, Ti}(origin::Tuple{T, T},
                                         antipode::Tuple{T, T},
                                         n::Tuple{Ti, Ti},
                                         locs::Array{Tuple{T, T}, 2},
                                         distances::Array{T, 2},
                                         steps::Tuple{T, T}) where {T<:Real, Ti<:Integer}
        all(n .== size(locs)) ||
            throw(DimensionMismatch("Locations are wrong dimension."))
        all(prod(n) .== size(distances)) ||
            throw(DimensionMismatch("Distances are wrong dimension."))
        new(origin, antipode, n, locs, distances, steps)
    end
end
"""
    GriddedFisheryDomain(origin::Tuple{T, T},
                         antipode::Tuple{T, T},
                         n::Tuple{Tn, Tn}) where {T<:Real, Tn<:Integer}

Construct a rectangular GriddedFisheryDomain with near and far corners
origin and antipode respecitively, and n grid cells in each direction.
"""
function GriddedFisheryDomain(origin::Tuple{T, T},
                              antipode::Tuple{T, T},
                              n::Tuple{Tn, Tn}) where {T <: Real, Tn <: Integer}
    indices = Base.product(1:n[1], 1:n[2])
    ## Set grid up so that locations are at center of each grid cell
    domain_size = abs.(antipode .- origin)
    steps = domain_size ./ n
    offset = steps ./ 2
    locranges = LinRange.(origin .+ offset, antipode .- offset, n)
    locs = collect(Base.product(locranges...))

    distances = calculate_distances(locs)

    GriddedFisheryDomain{T, Tn}(origin, antipode, n, locs, distances, steps)
end
## Version where `n` is equal along each axis
function GriddedFisheryDomain(origin::Tuple{T, T},
                              antipode::Tuple{T, T},
                              n::Tn) where {T <: Real, Tn <: Integer}
    GriddedFisheryDomain(origin, antipode, (n, n))
end
## Default to 100×100 grid with 100 cells in each direction
function GriddedFisheryDomain()
    GriddedFisheryDomain((0.0, 0.0), (100.0, 100.0), (100, 100))
end

size(Ω::GriddedFisheryDomain) = Ω.n
length(Ω::GriddedFisheryDomain) = prod(Ω.n)
eachindex(Ω::GriddedFisheryDomain) = eachindex(Ω.locs)

"""
    getindex(Ω::GriddedFisheryDomain, idx::Integer)
    getindex(Ω::GriddedFisheryDomain, idx::Integer, jdx::Integer)

Return the coordinates of the center of the `idx` cell, linearly indexed, or the
(idx, jdx) cell in Cartesian indices.
 """
function getindex(Ω::GriddedFisheryDomain, idx::Integer)
    getindex(Ω.locs, idx)
end

function getindex(Ω::GriddedFisheryDomain, idx::Integer, jdx::Integer)
    getindex(Ω, LinearIndices(Ω.n)[idx, jdx])
end

# Randomly sample locations from domain
function sample(rng, Ω::GriddedFisheryDomain)
    sample(rng, 1:length(Ω))
end

function sample(Ω::GriddedFisheryDomain)
    sample(Random.GLOBAL_RNG, Ω)
end

function sample(rng, Ω::GriddedFisheryDomain, E::Integer; replace = true)
    sample(rng, 1:length(Ω), E; replace = replace)
end

function sample(Ω::GriddedFisheryDomain, E::Integer; replace = true)
    sample(Random.GLOBAL_RNG, Ω, E; replace = replace)
end

# Weighted sampling from domain
function sample(rng::Random.AbstractRNG,
                Ω::GriddedFisheryDomain,
                w::StatsBase.AbstractWeights)
    sample(rng, 1:length(Ω), w)
end

function sample(Ω::GriddedFisheryDomain,
                w::StatsBase.AbstractWeights)
    sample(Random.GLOBAL_RNG, Ω, w)
end

function sample(rng::Random.AbstractRNG,
                Ω::GriddedFisheryDomain,
                w::StatsBase.AbstractWeights,
                E::Integer;
                replace = true)
    N = length(Ω)
    sample(rng, 1:N, w, E; replace = replace)
end

function sample(Ω::GriddedFisheryDomain,
                w::StatsBase.AbstractWeights,
                E::Integer;
                replace = true)
    sample(Random.GLOBAL_RNG, Ω, w, E; replace = replace)
end

"""
    calculate_distances(locs::Array{Tuple{T, T}, 2})

Calcluate distances between points where locs is a matrix of tuples. Calculates
upper triangle of distances, then make symmetric, returning a Matrix{T, 2}.
"""
function calculate_distances(locs::Array{Tuple{T, T}, 2}) where T <: Real
    N = prod(size(locs))
    distances = Array{T, 2}(undef, N, N)
    for j in 1:N
        for i in 1:(j - 1)
            distances[i, j] = hypot((locs[i] .- locs[j])...)
        end
        distances[j, j] = zero(T)
    end
    Matrix(Symmetric(distances, :U))
end

"""
    map_symm(f::F, A::Array)

Map a function `f` over the elements of a symmetric matrix `A` by calculating
the upper-triangular elements and mirroring, returning the full matrix.
"""
function map_symm(f::F, A::Array{T, 2}) where {F<:Function, T}
    I, J = size(A)
    res = similar(A)
    for j in 1:J, i in 1:j
        res[i, j] = f(A[i, j])
    end
    Symmetric(res, :U)
end


text:
using Test, GeometryTypes, CatmullClark

include("test1.jl")
include("test2.jl")


text:
using Distributions

function block_bootstrap_DGP(fm::FactorModel, factor_innovations::Array{Float64, 2}, block_size=10)
end
block_bootstrap_DGP(fm::FactorModel) = block_bootstrap_DGP(fm, apply(randn, size(fm.x)))

function parametric_bootstrap_DGP(fm::FactorModel, factor_innovations::Array{Float64, 2})
    # draw x from a parametric bootstrap with the parameters (estimated factors and loadings) given in fm
    # Note: this is a DGP for the factor equation (x = factors*loadings' + e where e are the innovations) only
    x = fm.factors*fm.loadings' + factor_innovations
end
parametric_bootstrap_DGP(fm::FactorModel) = parametric_bootstrap_DGP(fm, apply(randn, size(fm.x)))  # if no arguments are applied all innovations are assumed standard normal

# static factor models:

function residual_bootstrap(fm::FactorModel, B::Int, stat::Function)
    stats = Array(Float64, B)
    for b in 1:B
        resampled_x = fm.factors[:, 1:fm.number_of_factors]*fm.loadings[:, 1:fm.number_of_factors]' + fm.factor_residuals[rand(Distributions.DiscreteUniform(1, size(fm.x, 1)), size(fm.x, 1)), :]  # resample the residuals (by T index)
        stats[b] = stat(FactorModel(resampled_x, fm.number_of_factors, fm.number_of_factors_criterion, fm.factor_type, fm.targeted_predictors))
    end
    stats
end

function wild_bootstrap(fm::FactorModel, B::Int, stat::Function)  # resample residuals and multiply by random Variable with mean 0 and variance 1, do that B times and calculate statistic stat
    stats = Array(Float64, B)
    for b in 1:B
        resampled_x = fm.factors[:, 1:fm.number_of_factors]*fm.loadings[:, 1:fm.number_of_factors]' + apply(vcat, [fm.factor_residuals[t, :] .* randn() for t in 1:size(resampled_factor_residuals, 1)])
        #factor_residuals = fm.factor_residuals[rand(Distributions.DiscreteUniform(1, size(fm.x, 1)), size(fm.x, 1)), :]
        #resampled_x = fm.factors[:, 1:fm.number_of_factors]*fm.loadings[:, 1:fm.number_of_factors]' + apply(vcat, [resampled_factor_residuals[t, :] .* randn() for t in 1:size(resampled_factor_residuals, 1)])
        stats[b] = stat(FactorModel(resampled_x, fm.number_of_factors, fm.number_of_factors_criterion, fm.factor_type, fm.targeted_predictors))
    end
    stats
end

#y, x, f, lambdas, epsilon_x = factor_model_DGP(100, 60, 1; model="Breitung_Eickmeier_2011", b=2)
#fm = FactorModel(x, "ICp2")
#dfm = DynamicFactorModel((x, "ICp2"), 3)


text:
## Reading Logs
function readtb(logdir::String)
    hist = MVHistory()
    TensorBoardLogger.map_summaries(logdir) do tag, iter, val
        push!(hist, Symbol(tag), iter, val)
    end
    hist
end

function readtb(logdir::String, key)
    h = readtb(logdir)[key]
    h.iterations, h.values
end

# Convert various inputs to an arrow of directories for plotting
directories(input::String) = [input]
directories(input::T) where {T <: Solver} = [input.log.logger.logdir ]
directories(input::AbstractArray{Any}) = [i isa Solver ? i.log.logger.logdir : i for i in input]
directories(input::AbstractArray{T}) where {T <: String} = input
directories(input::AbstractArray{T}) where {T <: Solver} = [input[i].log.logger.logdir for i=1:length(input)]

## Learning Curve Metrics
# compute a percentile between a known range
percentile(frac, min, max) = frac*(max - min) + min

# Find the point in x that y exceeds val
function find_crossing(x, y, val)
    i = findfirst(y .>= val)
    isnothing(i) ? NaN : x[i]
end

# Exponential Smoothing
function smooth(v, weight = 0.6)
    N = length(v)
    smoothed = Array{Float64, 1}(undef, N)
    smoothed[1] = v[1]
    for i = 2:N
        smoothed[i] = smoothed[i-1] * weight + (1 - weight) * v[i]
    end
    smoothed
end

## Continual Learning Analysis
# Compiles the cumulative reward while learning a sequence of tasks
function cumulative_rewards(input, key=(i) -> Symbol("undiscounted_return/T$i"))
    dirs = directories(input)
    x, y, breaks, max_iter = [], [], [], 0
    for (dir, i) in zip(dirs, 1:length(dirs))
        hist = readtb(dir)[key(i)]
        push!(y, hist.values...)
        push!(x, (hist.iterations .+ max_iter)...)
        max_iter = maximum(x)
        push!(breaks, max_iter)
    end
    x, cumsum(y), breaks
end

# Compiles the performance on each task over time
function single_task_performances(input, key=(i) -> Symbol("undiscounted_return/T$i"))
    dirs = directories(input)
    res = Dict()
    breaks = []
    offset = 0
    for t=1:length(dirs) # Loop over each task
        for i=1:t # Loop up to each task
            try
            x, y = readtb(dirs[t], key(i)) # read the data
            if i==t
                res[i] = ([], []) # Empty data for task t
                push!(breaks, maximum(x))
            end

            push!(res[i][1], offset .+ x ...)
            push!(res[i][2], y...)
            catch end
        end
        offset = maximum(res[t][1])
    end
    res, cumsum(breaks)
end

function plot_cumulative_rewards(input; p=plot(), key=(i)->Symbol("undiscounted_return/T$i"), show_lines=false, kwargs...)
    x, y, lines = cumulative_rewards(input, key)

    plot!(p, x, y, title="Cumulative Reward", xlabel="Training Iteration", ylabel="Return"; kwargs...)

    if show_lines
        yrange = [0.9, 1.5].*[extrema(y)...]
        for l in lines
            plot!(p, [l,l], yrange, color=:black, label = "")
        end
    end
    p
end

function plot_jumpstart(input; p=plot(), key=(i)->Symbol("undiscounted_return/T$i"), kwargs...)
    dirs = directories(input)
    jumpstarts = [readtb(dir, key(i))[2][1] for (dir, i) in zip(dirs,1:length(dirs))]
    plot!(p, jumpstarts, xlabel = "Task Iteration", marker=true, markersize=3, ylabel="Return", title="Jumpstart Performance"; kwargs...)
end

function plot_peak_performance(input; p=plot(), key=(i)->Symbol("undiscounted_return/T$i"), smooth_weight=0.6, kwargs...)
    dirs = directories(input)
    jumpstarts = [maximum(smooth(readtb(dir, key(i))[2], smooth_weight)) for (dir, i) in zip(dirs,1:length(dirs))]
    plot!(p, jumpstarts, xlabel = "Task Iteration", marker=true, markersize=3, ylabel="Return", title="Peak Performance"; kwargs...)
end

function plot_steps_to_threshold(input, thresh; p=plot(), key=(i)->Symbol("undiscounted_return/T$i"), smooth_weight=0.6, kwargs...)
    dirs = directories(input)
    steps = []
    for (dir,i) in zip(dirs, 1:length(dirs))
        iters, vals = readtb(dir, key(i))
        push!(steps, find_crossing(iters, smooth(vals, smooth_weight), thresh)[1])
    end
    plot!(p, steps, xlabel = "Task Iteration", marker = true, markersize = 3, ylabel = "Iterations", title = "Steps to $thresh performance"; kwargs...)
end

function plot_forgetting(input; p=plot([plot() for _=1:length(input)]..., layout=(length(input),1)), key=(i)->Symbol("undiscounted_return/T$i"), smooth_weight=0.6, kwargs...)
    curves, lines = single_task_performances(input)
    for (k,v) in curves
        plot!(p[k], v[1], smooth(v[2], smooth_weight), xlabel="Training Iteration", ylabel="Return", xlims=(0,Inf); kwargs...)
    end
    p
end

function plot_learning(input;
        title = "",
        ylabel = "Undiscounted Return",
        xlabel = "Training Steps",
        values = :undiscounted_return,
        labels = :default,
        legend = :bottomright,
        font = :palatino,
        p = plot(),
        vertical_lines = [],
        vline_range = (0, 1),
        thick_every = 1
    )
    dirs = directories(input)

    N = length(dirs)
    values isa Symbol && (values = fill(values, N))
    if labels == :default
        labels = N == 1 ? [""] : ["Task $i" for i=1:N]
    end

    # Plot the vertical lines (usually for multitask learning or to designate a point on the curve)
    for i = 1:length(vertical_lines)
        plot!(p, [vertical_lines[i], vertical_lines[i]], [vline_range...], color=:black, linewidth = i % thick_every == 0 ? 3 : 1, label = "")
    end

    # Plot the learning curves
    plot!(p, ylabel = ylabel, xlabel = xlabel, legend = legend, title = title, fontfamily = font)
    for i in 1:length(dirs)
        x, y = readtb(dirs[i], values[i])
        plot!(p, x, y, alpha = 0.3, color=i, label = "")
        plot!(p, x, smooth(y), color=i, label = labels[i], linewidth =2 )
    end
    p
end

## Visualization
function episode_frames(mdp, policy; Neps=1, max_steps=1000, use_obs=false, render_kwargs::NamedTuple=(;), S=state_space(mdp))
    frames = []
    for i = 1:Neps
        s = rand(initialstate(mdp))
        o = mdp isa POMDP ? rand(initialobs(mdp, s)) : convert_s(AbstractArray, s, mdp)
        step = 0
        while !isterminal(mdp, s) && step < max_steps
            step += 1
            a = action(policy, tovec(o,S))
            size(a) == (1,) && (a=a[1])
            if mdp isa POMDP
                s, o, _ = gen(mdp, s, a)
            else
                s, _ = gen(mdp, s, a)
                o = convert_s(AbstractArray, s, mdp)
            end
            use_obs ? push!(frames, o') : push!(frames, render(mdp, s, a; render_kwargs...))
        end
    end
    frames
end

function gif(mdp, data::ExperienceBuffer, filename; convert_s = (s)->s, convert_a=(a)->a, render_kwargs::NamedTuple=(;), kwargs...)
    frames = [render(mdp, convert_s(bslice(data[:s], i)), convert_a(bslice(data[:a], i)); render_kwargs...) for i=1:length(data)]
    gif(frames, filename; kwargs...)
end

gif(mdp, policy, filename; fps=15, Neps=1, max_steps=1000, use_obs=false, S=state_space(mdp), render_kwargs::NamedTuple=(;)) = gif(episode_frames(mdp, policy, Neps=Neps, max_steps=max_steps, use_obs=use_obs, render_kwargs=render_kwargs, S=S), filename, fps=fps)

function gif(frames, filename; fps=15)
    save(filename, cat(frames..., dims=3), fps=fps)
end


text:
module SimpleCrop

using Cropbox

include("plant.jl")
include("soilwater.jl")
include("weather.jl")

@system Model(Plant, SoilWater, Weather, Controller)

export Model

end


text:
using DataFrames
using CSV
include("harp_reg_addr.jl")

function read_harp_bin(file_name::String)
    harp_file = open(file_name,"r")
    Message   = UInt8[]
    Timestamp = Float64[]
    Addresses = UInt8[]
    Payloads  = []
    Types     = UInt8[]
    while !eof(harp_file)
        try
            message_byte = read(harp_file,UInt8)
            length       = read(harp_file,UInt8)
            address      = read(harp_file,UInt8)
            port         = read(harp_file,UInt8)
            payload_type = read(harp_file,UInt8)
            seconds      = read(harp_file,UInt32)
            micro        = read(harp_file,UInt16)
            time         = seconds + micro * 32e-6
            payload_size = length - 0x0a
            payload_data = read(harp_file,payload_size)
            payload      = read(IOBuffer(payload_data),payloadtypes[payload_type  & ~0x10])
            checksum= read(harp_file,UInt8)
            push!(Message,message_byte)
            push!(Timestamp,time)
            push!(Addresses,address)
            push!(Payloads,payload)
            push!(Types,payload_type * ~0x10)
        catch e

            if isa(e, EOFError)
                mem_file = string(file_name,"memdump.txt")
                # Add code to dump current memory state to this mem-file in hex
                println("EOFError in $file_name - processing continuing.  Manually check.")
                break
            else
                ErrorException("Unhandled error when processing $file_name .  Program terminating")
            end
        end
    end
    return Message, Timestamp, Addresses,Payloads, Types
end

function read_message(x::UInt16)
    N = count_ones(x)
    inds = Vector{Int}(undef, N)
    if N == 0
        return inds
    end
    k = trailing_zeros(x)
    x >>= k + 1
    i = N - 1
    inds[N] = n = 16 - k
    while i >= 1
        (x, r) = divrem(x, 0x2)
        n -= 1
        if r == 1
            inds[i] = n
            i -= 1
        end
    end
    return sort(16 .- inds)
end

function common_elements(ind_array::Vector{Int},sub_array::Vector{Int})
    common_ind = Array{Union{Nothing,Int}}(nothing,length(sub_array))
    for (ind,element) in enumerate(sub_array)
        common_ind[ind] = findfirst(element .== ind_array)
    end
    return common_ind
end
function track_pwm(Message,Timestamp,Addresses,Payloads,Types,writes)
    on_pwm = findall(Addresses[writes] .=== 0x44)
    off_pwm= findall(Addresses[writes] .=== 0x45)
    for on_command in on_pwm
        on_target = Payloads[on_pwm[1]]
    end
end

function track_state(Message,Timestamp,Addresses,Payloads,Types)
    num_states = length(registerbits_A)
    state_bits = sort(collect(keys(registerbits_A)))
    state_names= Vector{String}()
    for ind in state_bits
        push!(state_names,registerbits_A[ind])
    end
    reads  = Message .=== 0x01
    writes = Message .=== 0x02
    events = Message .=== 0x03
    num_writes = count(writes)
    # Construct events
    event_types = unique(Types[events])
    event_dictionary = Dict{UInt8,DataFrame}()
    for cur_type in event_types
        type_index = (Types .=== cur_type) .& events
        event_address = unique(Addresses[type_index])
        for cur_address in event_address
            event_index = (Addresses .=== cur_address) .& type_index
            merge!(event_dictionary,Dict([(cur_address,DataFrame(Timestamp = Timestamp[event_index], Payload = Payloads[event_index]))]))
        end
    end
    write_events        = Array{Union{Nothing,Bool,Float64,String},2}(nothing,num_writes,num_states+1)
    write_sets          = writes .& (Addresses .=== 0x22)
    write_set_events    = Array{Union{Nothing,Bool},2}(nothing,count(write_sets),num_states)
    write_clears        = writes .& (Addresses .=== 0x23)
    write_clear_events  = Array{Union{Nothing,Bool},2}(nothing,count(write_clears),num_states)
    write_toggles       = writes .& (Addresses .=== 0x24)
    write_toggle_events = Array{Union{Nothing,String},2}(nothing,count(write_toggles),num_states)
    set_inds    = findall(Addresses[writes] .=== 0x22)
    clear_inds  = findall(Addresses[writes] .=== 0x23)
    toggle_inds = findall(Addresses[writes] .=== 0x24)
    for (ind,write_ind) in enumerate(findall(write_sets))
        cur_payload = read_message(Payloads[write_ind])
        set_registers = common_elements(state_bits,cur_payload)
        write_set_events[ind,set_registers] .= true
    end
    for (ind,write_ind) in enumerate(findall(write_clears))
        cur_payload = read_message(Payloads[write_ind])
        clear_registers = common_elements(state_bits,cur_payload)
        write_clear_events[ind,clear_registers] .= false
    end
    for (ind,write_ind) in enumerate(findall(write_toggles))
        cur_payload = read_message(Payloads[write_ind])
        toggle_registers = common_elements(state_bits,cur_payload)
        write_toggle_events[ind,toggle_registers] .= "TOGG"
    end
    write_events[set_inds,2:end]   = write_set_events
    write_events[clear_inds,2:end] = write_clear_events
    write_events[toggle_inds,2:end]= write_toggle_events

    write_events[:,1] = Timestamp[writes]

    # Now delete superflous columnns
    col_with_data  = dropdims(sum(write_events .=== nothing,dims=1) .!= size(write_events)[1],dims=1)
    write_events   = write_events[:,col_with_data]
    write_headings = vcat("Times",state_names)
    write_headings = write_headings[col_with_data]
    # Fill in the extra state
    for row_ind = 2:size(write_events)[1]
        propagate_state = (write_events[row_ind,:] .=== nothing) .& (write_events[row_ind-1,:] .!== "TOGG")
        write_events[row_ind,propagate_state] = write_events[row_ind-1,propagate_state]
    end
    write_data = DataFrame(write_events,write_headings)
    return event_dictionary, write_data
end

function sink_data(event_dictionary::Dict{UInt8, DataFrame},write_data::DataFrame,file_name::String)
    for col in eachcol(write_data)
        replace!(col,nothing=>NaN)
    end
    base_title = replace(file_name,".bin"=>"")
    for (key, value) in event_dictionary
        title = string(base_title,"_event_data_",key,".csv")
        CSV.write(title,value)
    end
    title = string(base_title,"_write_data.csv")
    CSV.write(title,write_data)
end

function sink_folder(event_dictionary::Dict{UInt8,DataFrame},write_data::DataFrame,file_name::String)
    for col in eachcol(write_data)
        replace!(col,nothing=>NaN)
    end

    base_title = replace(file_name,".bin"=>"")
    mkdir(base_title)
    for (key, value) in event_dictionary
        title = joinpath(base_title,string("event_data_",key,".csv"))
        CSV.write(title,value)
    end
    title = joinpath(base_title,"write_data.csv")
    CSV.write(title,write_data)
end

max_stars_repo_name: Peter-Vincent/Harpist

max_issues_repo_name: Peter-Vincent/Harpist

max_forks_repo_name: Peter-Vincent/Harpist


text:
include("radCalcs.jl")
using QuadGK

"""
Function for calculating the OMode group delay.

Inputs:
        freq::Vector{Float64}         -- Measured frequencies of the reflectometer (GHz)
        coeffs_fit::Vector{Float64}   -- The 6 fit coefficients for the 5th order polynomial for plasma density in units of 1e19 (per cubic meter)
                                            n(r) = 1e19*(coeffs_fit[1]*r^5+coeffs_fit[2]*r^4+coeffs_fit[3]*r^3+coeffs_fit[4]*r^2+coeffs_fit[5]*r^1+coeffs_fit[6])
                                            where r=radius into the plasma from the plasma facing wall.
        calibration::Vector{Float64}  -- Calibration length in meters (non-plasma group delay)
        rad::Vector{Float64}          -- Radial location of cutoffs are each frequency in freq (meters from first wall).

Output:
        ophiw   -- the group delay at each frequency in freq
"""
function dphidw_omode(freq, coeffs_fit::Vector{Float64},calibration::Vector{Float64},rad::Vector{Float64})
    ophi = Array{Float64}(undef,length(freq))
    for j in 1:1:length(freq)
        ophi[j] = calibration[j] / 3e8 * (freq[j]*2*pi*1e9) + 2*(freq[j]*2*pi*1e9) / 3e8 * quadgk(x -> refractive_omode(x;coeffs = coeffs_fit, omega = freq[j]*2*pi*1e9),0,(rad[j]),rtol=1e-12)[1]
    end
    ophiw = Array{Float64}(undef,length(freq))
    dw = (freq[2]-freq[1])*2*pi*1e9
    ophiw[1] = (ophi[2]-ophi[1])/dw
    ophiw[length(freq)] = (ophi[length(freq)]-ophi[length(freq)-1])/dw
    for j in 2:1:length(ophiw)-1
        ophiw[j] = (ophi[j+1]-ophi[j-1])/(2*dw)
    end

    ophiw
end

"""
Function for calculating the XMode group delay.

Inputs:
        freq                          -- Measured frequencies of the reflectometer (GHz)
        coeffs_fit::Array{Float64}    -- The 6 fit coefficients for the 5th order polynomial for plasma density in units of 1e19 (per cubic meter)
                                            n(r) = 1e19*(coeffs_fit[1]*r^5+coeffs_fit[2]*r^4+coeffs_fit[3]*r^3+coeffs_fit[4]*r^2+coeffs_fit[5]*r^1+coeffs_fit[6])
                                            where r=radius into the plasma from the plasma facing wall.
        calibration::Array{Float64}   -- Calibration length in meters (non-plasma group delay)
        rad::Array{Float64}           -- Radial location of cutoffs are each frequency in freq (meters from first wall).
        Bmag::Number                  -- The magnitude of the magnetic field at major radius startingRad (Tesla)
        startingRad::Number           -- The major radial measurement location of the magnetic field (meters).

Output:
        ophiw -- the group delay at each frequency in freq
"""
function dphidw_xmode(freq, coeffs_fit::Vector{Float64}, calibration::Array{Float64}, Bmag::Number, startingRad::Number,rad::Array{Float64})
    ophi = Array{Float64}(undef,length(freq))
    for j in 1:1:length(freq)
        ophi[j] = calibration[j] / 3e8 * (freq[j]*2*pi*1e9) + 2*(freq[j]*2*pi*1e9) / 3e8 * quadgk(x -> refractive_xmode(x;coeffs = coeffs_fit, omega = freq[j]*2*pi*1e9, b0 = Bmag, r0 = startingRad),0,(rad[j]),rtol=1e-12)[1]
    end
    ophiw = Array{Float64}(undef,length(freq))
    dw = (freq[2]-freq[1])*2*pi*1e9
    ophiw[1] = (ophi[2]-ophi[1])/dw
    ophiw[length(freq)] = (ophi[length(freq)]-ophi[length(freq)-1])/dw
    for j in 2:1:length(ophiw)-1
        ophiw[j] = (ophi[j+1]-ophi[j-1])/(2*dw)
    end

    ophiw
end

"""
Function for calculating the OMode index of refraction

Inputs:
        x::Float64                    -- the radial location in meters from the first wall
        coeffs_fit::Array{Float64}    -- The 6 fit coefficients for the 5th order polynomial for plasma density in units of 1e19 (per cubic meter)
                                            n(r) = 1e19*(coeffs_fit[1]*r^5+coeffs_fit[2]*r^4+coeffs_fit[3]*r^3+coeffs_fit[4]*r^2+coeffs_fit[5]*r^1+coeffs_fit[6])
                                            where r=radius into the plasma from the plasma facing wall.
         omega::Float64               -- the excited frequency (radians per second)

Output:
        -- the OMode index of refraction
"""
function refractive_omode(x::Float64; coeffs::Array{Float64}, omega::Float64)
    q = 1.6022e-19
    me = 9.109e-31
    e0 = 8.854e-12
    density = 1e19.*(coeffs[1].*(x).^5 .+coeffs[2].*(x).^4 .+coeffs[3].*(x).^3 .+coeffs[4].*(x).^2 .+coeffs[5].*(x).+coeffs[6])
    wpe2 = (q)^2 * density / (e0*me)
    sqrt(max(0.0,1.0 .-wpe2./omega^2))
end

"""
Function for calculating the XMode index of refraction

Inputs:
        x::Float64                    -- the radial location in meters from the first wall
        coeffs_fit::Array{Float64}    -- The 6 fit coefficients for the 5th order polynomial for plasma density in units of 1e19 (per cubic meter)
                                            n(r) = 1e19*(coeffs_fit[1]*r^5+coeffs_fit[2]*r^4+coeffs_fit[3]*r^3+coeffs_fit[4]*r^2+coeffs_fit[5]*r^1+coeffs_fit[6])
                                            where r=radius into the plasma from the plasma facing wall.
         omega::Float64               -- the excited frequency (radians per second)
         b0::Number                   -- The magnitude of the magnetic field at major radius startingRad (Tesla)
         r0::Number                   -- The major radial measurement location of the magnetic field (meters).

Output:
        -- the XMode index of refraction
"""
function refractive_xmode(x::Float64; coeffs::Array{Float64}, omega::Float64, b0::Number, r0::Number)
    q = 1.6022e-19
    me = 9.109e-31
    e0 = 8.854e-12
    density = 1e19.*(coeffs[1].*(x).^5 .+coeffs[2].*(x).^4 .+coeffs[3].*(x).^3 .+coeffs[4].*(x).^2 .+coeffs[5].*(x).+coeffs[6])
    wpe2 = (q)^2 * density / (e0*me)
    wce = b0 / (r0+x) * q / me
    sqrt(max(0.0,1.0 .-wpe2./omega^2 * (omega^2-wpe2)/(omega^2-(wpe2+wce^2))))
end


text:
struct Verbosity{T}
end

result_color(cs::ReportColorScheme, ::Any, ::Any) = :default
result_color(cs::ReportColorScheme, ::Test.Pass, ::Verbosity{2}) = cs.color_pass
result_color(cs::ReportColorScheme, ::Test.Fail, ::Any) = cs.color_fail
result_color(cs::ReportColorScheme, ::Test.Error, ::Any) = cs.color_error
result_color(cs::ReportColorScheme, ::Test.Broken, ::Verbosity{2}) = cs.color_broken
result_color(cs::ReportColorScheme, ::ReturnValue, ::Verbosity{2}) = cs.color_return
result_color(cs::ReportColorScheme, ::FailExplanation, ::Any) = cs.color_fail

result_show(::Test.Pass, ::Verbosity{1}) = "."
result_show(::Test.Pass, ::Verbosity{2}) = "PASS"
result_show(::Test.Broken, ::Verbosity{1}) = "B"
result_show(::Test.Broken, ::Verbosity{2}) = "BROKEN"
result_show(::ReturnValue, ::Verbosity{1}) = "*"
# Since the `show` method for the result type will probably be defined in the test file,
# we need to use `invokelatest` here for it to be picked up.
result_show(result::ReturnValue, ::Verbosity{2}) = Base.invokelatest(string, result.value)
result_show(::Test.Fail, ::Verbosity{1}) = "F"
result_show(::Test.Fail, ::Verbosity{2}) = "FAIL"
result_show(::Test.Error, ::Verbosity{1}) = "E"
result_show(::Test.Error, ::Verbosity{2}) = "ERROR"
result_show(::FailExplanation, ::Verbosity{1}) = "F"
result_show(::FailExplanation, ::Verbosity{2}) = "FAIL"
result_show(::JuteTestSet, ::Verbosity) = ""

function Base.show(io::IO, fe::FailExplanation)
    printstyled(io, "Test Failed\n"; color=:red, bold=true)
    descr = replace(fe.description, r"^"m => "  ")
    print(io, descr)
end

struct ExecutedTestcase
    tcinfo :: TestcaseInfo
    labels :: Array{String, 1}
    outcome :: TestcaseOutcome
end

mutable struct ProgressReporter
    verbosity :: Int
    just_started :: Bool
    current_group :: Array{String, 1}
    doctest :: Bool
    start_time :: UInt64
    visited_groups :: Set{Array{String, 1}}
    testcases :: Array{ExecutedTestcase, 1}

    function ProgressReporter(verbosity, doctest)
        new(verbosity, true, String[], doctest, 0, Set{Array{String, 1}}(), ExecutedTestcase[])
    end
end

function common_elems_num(l1, l2)
    res = 0
    for i in 1:min(length(l1), length(l2))
        if l1[i] != l2[i]
            return res
        end
        res = i
    end
    return res
end

function _progress_start_testgroup!(progress::ProgressReporter, tcinfo::TestcaseInfo)

    verbosity = progress.verbosity
    path = tcinfo.path

    if verbosity == 1 && !progress.just_started
        println()
    end

    if length(path) > 0
        path_repeated = path in progress.visited_groups

        if path_repeated && verbosity == 1
            cn = 0
        else
            cn = common_elems_num(progress.current_group, path)
        end

        for i in cn+1:length(path)
            postfix = verbosity == 1 ? (path_repeated ? " (cont.):" : ":") : "/"
            print("  " ^ (i - 1), path[i], postfix)
            if verbosity == 1
                if i != length(path)
                    print("\n")
                else
                    print(" ")
                end
            elseif verbosity == 2
                print("\n")
            end
        end
    end

end

function progress_start_testcase!(progress::ProgressReporter, tcinfo::TestcaseInfo, labels)

    progress.just_started = false

    path, name = path_pair(tcinfo)
    verbosity = progress.verbosity

    if verbosity > 0 && path != progress.current_group
        _progress_start_testgroup!(progress, tcinfo)
        progress.current_group = path
        push!(progress.visited_groups, path)
    end

    if progress.verbosity >= 2
        tctag = tag_string(tcinfo, labels)
        path, name = path_pair(tcinfo)
        print("  " ^ length(path), tctag, " ")
    end
end

function progress_finish_testcase!(
        progress::ProgressReporter, tcinfo::TestcaseInfo, labels,
        outcome::TestcaseOutcome, color_scheme::ReportColorScheme)

    push!(progress.testcases, ExecutedTestcase(tcinfo, labels, outcome))

    verbosity = progress.verbosity
    if verbosity == 1
        for result in outcome.results
            printstyled(
                result_show(result, Verbosity{verbosity}()),
                color=result_color(color_scheme, result, Verbosity{verbosity}()))
        end
    elseif verbosity >= 2
        if progress.doctest
            elapsed_time = "[...] ms"
        else
            elapsed_time = pprint_time(outcome.elapsed_time)
        end

        print("($elapsed_time)")

        for result in outcome.results
            result_str = result_show(result, Verbosity{verbosity}())
            if result_str != ""
                printstyled(
                    " [$result_str]",
                    color=result_color(color_scheme, result, Verbosity{verbosity}()))
            end
        end
        println()
    end
end

function progress_start!(progress::ProgressReporter)
    if progress.verbosity > 0

        if progress.doctest
            platform = "[...], Julia [...]"
            jute_version = "[...]"
        else
            kernel = Sys.iswindows() ? "Windows" : Sys.isapple() ? "macOS" : Sys.KERNEL
            platform = "$kernel ($(Sys.MACHINE)), Julia $(string(VERSION))"
            if VERSION < v"1.4"
                jute_version = get(Pkg.installed(), "Jute", "<unknown>")
            else
                # The new way, as per https://github.com/JuliaLang/Pkg.jl/issues/385
                # Works sometimes. Or, at least, allows the tests to pass.
                jute_version = Pkg.project().version
            end
        end

        println("Platform: $platform, Jute $jute_version")
        println("-" ^ 80)
    end

    progress.start_time = time_ns()
end

function _custom_results_present(progress::ProgressReporter)
    for etc in progress.testcases
        if any(isa(result, ReturnValue) for result in etc.outcome.results)
            return true
        end
    end
    false
end

function _print_custom_results(progress::ProgressReporter, color_scheme::ReportColorScheme)
    println("-" ^ 80)
    pr2 = ProgressReporter(2, false)
    for etc in progress.testcases
        return_values = [result for result in etc.outcome.results if isa(result, ReturnValue)]
        if !isempty(return_values)
            filtered_outcome = TestcaseOutcome(
                return_values, etc.outcome.elapsed_time, etc.outcome.output)
            progress_start_testcase!(pr2, etc.tcinfo, etc.labels)
            progress_finish_testcase!(pr2, etc.tcinfo, etc.labels, filtered_outcome, color_scheme)
        end
    end
end

function _print_statistics(progress::ProgressReporter, full_time)

    all_results = mapreduce(etc -> etc.outcome.results, vcat, progress.testcases, init=[])
    stats = Dict(
        key => length(filter(result -> isa(result, tp), all_results))
        for (key, tp) in [
            (:pass, Union{Test.Pass, ReturnValue, Test.Broken}),
            (:fail, Union{Test.Fail, FailExplanation}),
            (:error, Test.Error)])

    if progress.doctest
        full_time_str = "[...] s"
        full_test_time_str = "[...] s"
    else
        full_test_time = mapreduce(etc -> etc.outcome.elapsed_time, +, progress.testcases)
        full_time_str = pprint_time(full_time, meaningful_digits=3)
        full_test_time_str = pprint_time(full_test_time, meaningful_digits=3)
    end

    println("-" ^ 80)
    println(
        "$(stats[:pass]) tests passed, " *
        "$(stats[:fail]) failed, " *
        "$(stats[:error]) errored " *
        "in $full_time_str (total test time $full_test_time_str)")
end

function _print_failures(progress::ProgressReporter)
    for etc in progress.testcases
        if is_failed(etc.outcome)
            println("=" ^ 80)
            println(tag_string(etc.tcinfo, etc.labels; full=true))

            if length(etc.outcome.output) > 0
                println("Captured output:")
                println(etc.outcome.output)
            end

            for result in etc.outcome.results
                if is_failed(result)
                    println(result)
                end
            end
        end
    end
end

function progress_finish!(progress::ProgressReporter, color_scheme::ReportColorScheme)

    full_time = (time_ns() - progress.start_time) / 1e9

    if progress.verbosity == 1
        # Results are displayed with `print()`, so a line break in the end is needed.
        println()
    end

    if progress.verbosity == 1 && _custom_results_present(progress)
        _print_custom_results(progress, color_scheme)
    end

    if progress.verbosity >= 1
        _print_statistics(progress, full_time)
    end

    _print_failures(progress)
end


text:
# ---
# title: 771. Jewels and Stones
# id: problem771
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Hash Table
# link: <https://leetcode.com/problems/jewels-and-stones/description/>
# hidden: true
# ---
#
# You're given strings `J` representing the types of stones that are jewels, and
# `S` representing the stones you have.  Each character in `S` is a type of
# stone you have.  You want to know how many of the stones you have are also
# jewels.
#
# The letters in `J` are guaranteed distinct, and all characters in `J` and `S`
# are letters. Letters are case sensitive, so `"a"` is considered a different
# type of stone from `"A"`.
#
# **Example 1:**
#
#
#
#     Input: J = "aA", S = "aAAbbbb"
#     Output: 3
#
#
# **Example 2:**
#
#
#
#     Input: J = "z", S = "ZZ"
#     Output: 0
#
#
# **Note:**
#
#   * `S` and `J` will consist of letters and have length at most 50.
#   * The characters in `J` are distinct.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end