"""Module Sario
"""

__precompile__(true)

module Sario

include("./demrsc.jl")

export load, save, DemRsc, load_stack, load_hdf5_stack, load_mask, load_geolist_from_h5,
       load_intlist_from_h5, load_dem_from_h5, get_file_ext, find_files, save_hdf5_stack

# TODO: python transfers:
# const LOAD_IN_PYTHON = vcat(IMAGE_EXTS, [".geojson", ".npy"])
# GeoJSON.geo2dict(GeoJSON.parsefile("../../bbox2.geojson"))["coordinates"]

import Glob
import JSON
using HDF5
using Dates
using Parameters

const DATE_FMT = "yyyymmdd"
const SENTINEL_EXTS = [".geo", ".cc", ".int", ".amp", ".unw", ".unwflat"]
const COMPLEX_EXTS = [".int", ".slc", ".geo", ".cc", ".unw", ".unwflat", ".mlc", ".grd"]
const REAL_EXTS = [".amp", ".cor", ".mlc", ".grd"]
const ELEVATION_EXTS = [".dem", ".hgt"]
# These file types are not simple complex matrices: see load_stacked_img for detail
# .unwflat are same as .unw, but with a linear ramp removed
const STACKED_EXTS = [".cc", ".unw", ".unwflat"]
const BOOL_EXTS = [".mask"]

const IMAGE_EXTS = [".png", ".tif", ".tiff", ".jpg"]
const NOT_IMPLEMENTED = vcat(IMAGE_EXTS, [".geojson", ".npy"])

# For HDF5 saving
const DEM_RSC_DSET = "dem_rsc"
const STACK_DSET = "stack"
const GEOLIST_DSET = "geo_dates"
const INTLIST_DSET = "int_dates"

RSC_KEY_TYPES = [
    ("width", Int),
    ("file_length", Int),
    ("x_first", Float64),
    ("y_first", Float64),
    ("x_step", Float64),
    ("y_step", Float64),
    ("x_unit", String),
    ("y_unit", String),
    ("z_offset", Int),
    ("z_scale", Int),
    ("projection", String),
]
RSC_KEYS = [tup[1] for tup in RSC_KEY_TYPES]

find_files(ext, directory = ".") = sort(Glob.glob("*" * ext, directory))

"""Extracts the file extension, including the "." (e.g.: .slc)"""
get_file_ext(filename::AbstractString) = splitext(filename)[end]

_is_h5(filename) = get_file_ext(filename) in (".h5", ".hdf5")

import Base.size
function Base.size(h5file::String, dset::String)
    h5open(h5file) do f
        return size(f[dset])
    end
end

"""
    load(filename::AbstractString;
         rsc_file::Union{AbstractString, Nothing}=nothing,
         looks::Tuple{Int, Int}=(1, 1), do_permute=true,
         dset_name::AbstractString="", return_amp::Bool=false)

Main entry point to loading files.

Examines file type of `filename` and runs appropriate load function.

"""
function load(filename::AbstractString;
              rsc_file::Union{AbstractString,Nothing} = nothing,
              looks::Tuple{Int,Int} = (1, 1), do_permute = true,
              dset_name::AbstractString = "", return_amp::Bool = false)
    ext = get_file_ext(filename)

    # For now, just pass through unimplemented extensions to Python
    if ext in NOT_IMPLEMENTED
        error("$ext is not yet implemented")
    elseif ext == ".rsc"
        return take_looks(_get_rsc_data(filename, filename), looks...)
    elseif _is_h5(filename)
        return take_looks(_load_hdf5(filename, dset_name; do_permute = do_permute), looks...)
    end

    # Sentinel files should have .rsc file: check for dem.rsc, or elevation.rsc
    data_type, demrsc, num_rows, num_cols = _file_info(filename, rsc_file)
    # If the filesize doesn't match the .rsc, error
    _check_filesize(filename, data_type, (num_rows, num_cols))

    if ext in ELEVATION_EXTS
        return take_looks(load_elevation(filename, demrsc, do_permute = do_permute), looks...)
    elseif ext in STACKED_EXTS
        return take_looks(load_stacked_img(filename, demrsc, do_permute = do_permute, return_amp = return_amp),
                          looks...)
    elseif ext in BOOL_EXTS
        return take_looks(load_bool(filename, demrsc, do_permute = do_permute), looks...)
    else
        return take_looks(load_complex(filename, demrsc, do_permute = do_permute), looks...)
    # having demrsc implies that this is not a UAVSAR file, so is complex
    # TODO: haven"t transferred over UAVSAR functions, so no load_real yet
    end

end

function _check_filesize(fname, data_type, dsize)
    fsize = Int(filesize(fname) / sizeof(data_type))
    # The stacked files have 2x the number of pixels of Floats- needs to be scaled to match dem.rsc
    fsize = get_file_ext(fname) in STACKED_EXTS ? fsize / 2 : fsize
    @assert prod(dsize) == fsize ".rsc size $dsize does not match $fname size $fsize"
end

"""Load one element of a file on disk (avoid reading in all of huge file"""
# TODO: Load a chunk of a file now?
function load(filename::AbstractString, row_col::Tuple{Int,Int}; rsc_file::Union{AbstractString,Nothing} = nothing)
    data_type, demrsc, num_rows, num_cols = _file_info(filename, rsc_file)

    row, col = row_col
    _check_bounds(row, col, num_rows, num_cols)

    seek_pos = _get_seek_position(row, col, num_cols, data_type)

    open(filename) do f
        seek(f, seek_pos)
        # This read syntax loads single `data_type`
        return read(f, data_type)
    end
end

"""For single element reading in binary files, seek to the right row, col"""
_get_seek_position(row, col, num_cols, data_type) = sizeof(data_type) * ((col - 1) + (num_cols * (row - 1)) )

function _check_bounds(row, col, num_rows, num_cols)
    if row < 1 || col < 1 || row > num_rows || col > num_cols
        throw(DomainError((row, col), " out of bounds for $filename of size ($num_rows, $num_cols)"))
    end
end

function _file_info(filename, rsc_file)
    @assert isfile(filename) "$filename does not exist."

    data_type = _get_data_type(filename)
    demrsc = _get_rsc_data(filename, rsc_file)
    @unpack rows, cols = demrsc
    return data_type, demrsc, rows, cols
end

# NOTE: the subset will only work for sequential data (complex), not the stacked filetypes
"""Load subset of a file on disk using range"""
RangeTuple = Tuple{T,T} where {T <: OrdinalRange}
function load(filename::AbstractString, idxs::RangeTuple;
              rsc_file::Union{AbstractString,Nothing} = nothing,  do_permute = true)
    data_type, demrsc, num_rows, num_cols = _file_info(filename, rsc_file)

    # Note: reversing these since julia uses column-major
    cols, rows = idxs
    _check_bounds(rows.start, cols.start, num_rows, num_cols)

    outrows = rows.stop  - rows.start + 1
    outcols = cols.stop  - cols.start + 1
    out = zeros(data_type, (outrows, outcols))
    buf = zeros(data_type, outcols)

    open(filename) do f
        for (idx, r) in enumerate(rows)
            seek_pos = _get_seek_position(r, cols.start, num_cols, data_type)
            seek(f, seek_pos)
            # This read syntax loads single `data_type`
            # read!(out[r, :], f, data_type)
            read!(f, buf)
            out[:, idx] .= buf
        end
    end
    return do_permute ? permutedims(out) : out
end

function _load_hdf5(h5file::AbstractString, dset_name::AbstractString = ""; do_permute = true)
    h5open(h5file, "r") do f
        if isempty(dset_name)
            if length(names(f)) == 1 && typeof(f[names(f)[1]]) == HDF5Dataset
                dset_name = names(f)[1]
                println("Loading \"$dset_name\" from $h5file")
            else
                error("More than 1 dset exists in $h5file, provide name to `load`")
            end
        end
        order = ndims(f[dset_name]) == 3 ? (2, 1, 3) : (2, 1)  # Handle 3D permuting
        return do_permute ? permutedims(read(f, dset_name), order) : read(f, dset_name)
        k
    end
end

"""Handle getting the .rsc data from either an image filename, or .rsc filename"""
function _get_rsc_data(filename, rsc_file::Union{AbstractString,Nothing})
    ext = get_file_ext(filename)

    demrsc = nothing
    if !isnothing(rsc_file)
        demrsc = load_dem_rsc(rsc_file)
    elseif ext in vcat(SENTINEL_EXTS, ELEVATION_EXTS, BOOL_EXTS)
        rsc_file = find_rsc_file(filename)
        demrsc = load_dem_rsc(rsc_file)
    end

    return demrsc
end

function find_rsc_file(filename = nothing; directory = nothing, verbose = false)
    if !isnothing(filename)
        dir_parts = splitpath(filename)[1:end - 1]
        directory = isempty(dir_parts) ? "." : joinpath(dir_parts...)
    end
    # Should be just elevation.dem.rsc (for .geo folder) or dem.rsc (for igrams)
    possible_rscs = find_files("*.rsc", directory)
    if verbose
        println("Searching $directory for rsc files")
        println("Possible rsc files:")
        println(possible_rscs)
    end
    if length(possible_rscs) < 1
        println("No .rsc file found in $directory")
        return nothing
    elseif length(possible_rscs) > 1
        pathlist = splitpath(filename)
        if length(pathlist) == 1
            filepath, fileonly = ".", pathlist[1]
        else
            filepath, fileonly = joinpath(pathlist[1:end - 1]...), pathlist[end]
        end
        rsc_nopaths = [splitpath(r)[end] for r in possible_rscs]
        if any(startswith(r, fileonly) for r in rsc_nopaths)
            possible_rscs = [joinpath(filepath, r) for r in rsc_nopaths if startswith(r, fileonly)]
        else
            error("$filename has multiple .rsc files in its directory: $possible_rscs")
        end
    end
    return abspath(possible_rscs[1])
end

"""Convert the text file into the DemRsc struct
Starts with a dict to gather all fiields, then unpacks using keywords args"""
function load_dem_rsc(filename, kwargs...)::DemRsc
    output_data = Dict{Symbol,Any}()

    for line in readlines(filename)
        for (field, valtype) in RSC_KEY_TYPES
            if startswith(line, uppercase(field))
                val = valtype <: AbstractString ? split(line)[2] : parse(valtype, split(line)[2])
                output_data[Symbol(lowercase(field))] = val
            end
        end
    end
    return DemRsc(;output_data...)
end

function _get_data_type(filename)
    ext = get_file_ext(filename)
    if ext in ELEVATION_EXTS
        return Int16
    elseif ext in STACKED_EXTS
        return Float32
    elseif ext in BOOL_EXTS
        return Bool
    else
        return ComplexF32
    end
end

"""Loads a digital elevation map from either .hgt file or .dem

.hgt is the NASA SRTM files given. Documentation on format here:
https://dds.cr.usgs.gov/srtm/version2_1/Documentation/SRTM_Topo.pdf
Key point: Big-endian 2 byte (16-bit) integers

.dem is format used by Zebker geo-coded and ROI-PAC SAR software
Only difference is data is stored little-endian (like other SAR data)

Note on both formats: gaps in coverage are given by INT_MIN -32768,
so either manually set data(data == np.min(data)) = 0,
or something like
    data = clamp(data, -10000, Inf)
"""
function load_elevation(filename, demrsc::DemRsc; do_permute = true)
    rows, cols = (demrsc.file_length, demrsc.width)
    data = Array{Int16,2}(undef, (cols, rows))
    read!(filename, data)

    # # TODO: Verify that the min real value will be above -1000
    # min_valid = -10000
    # # Set NaN values to 0
    # @. data[data < min_valid] = 0
    return do_permute ? permutedims(data) : data
end

function load_complex(filename::AbstractString, demrsc::DemRsc; do_permute = true)
    return _load_bin_matrix(filename, demrsc, ComplexF32, do_permute)
end

function load_bool(filename::AbstractString, demrsc::DemRsc; do_permute = true)
    return _load_bin_matrix(filename, demrsc, Bool, do_permute)
end

function _load_bin_matrix(filename, demrsc::DemRsc, dtype, do_permute::Bool)
    @unpack rows, cols = demrsc
    # Note: must be saved in c/row-major order, so loading needs a transpose
    out = Array{dtype,2}(undef, (cols, rows))

    read!(filename, out)
    return do_permute ? permutedims(out) : out
end

"""load_stacked_img is for the weird "ALT_LINE_DATA" formatted images:
https://web.stanford.edu/group/radar/softwareandlinks/sw/snaphu/snaphu_man1.html#FILE%20FORMATS

Format is two stacked matrices:
[[first], [second]] where the first "cols" number of floats
are the first matrix, next "cols" are second, etc.

For .unw height files, the first is amplitude, second is phase (unwrapped)
For .cc correlation files, first is amp, second is correlation (0 to 1)
"""
function load_stacked_img(filename::AbstractString, demrsc::DemRsc; do_permute = true, return_amp = false)
    @unpack rows, cols = demrsc
    # Note: must be saved in c/row-major order, so loading needs a transpose
    # out_left usually is amplitude data
    # Usually we are interested in out_right
    #
    # First make container for all of data
    out = Array{Float32,2}(undef, (2cols, rows))
    read!(filename, out)

    out_amp = @view out[1:cols, :]
    out_data = @view out[cols + 1:end, :]

    if return_amp
        return do_permute ? permutedims(cat(out_amp, out_data, dims = 3), (2, 1, 3)) : cat(out_amp, out_data, dims = 3)
    else
        return do_permute ? permutedims(out_data) : out_data
    end
end

function load_geolist_from_h5(h5file::AbstractString; geolist_dset = GEOLIST_DSET)
    geo_strings = h5read(h5file, geolist_dset)
    return parse_geolist_strings(geo_strings)
end

function load_geolist_from_h5(h5file::AbstractString, dset::AbstractString; geolist_dset = GEOLIST_DSET)
    geo_strings = h5readattr(h5file, dset)[geolist_dset]
    return parse_geolist_strings(geo_strings)
end

function load_intlist_from_h5(h5file)
    # Note transpose, since it's stored as a N x 2 array
    # (which is flipped to 2 x N for julia
    int_strings = h5read(h5file, INTLIST_DSET)
    return parse_intlist_strings([Tuple(int_strings[:, ii]) for ii in 1:size(int_strings, 2)])
end

function load_intlist_from_h5(h5file, dset::AbstractString)
    int_strings = h5readattr(h5file, dset)[INTLIST_DSET]
    return parse_intlist_strings([Tuple(int_strings[:, ii]) for ii in 1:size(int_strings, 2)])
end

_parse(datestr) = Date(datestr, DATE_FMT)

function parse_intlist_strings(date_pairs::AbstractArray{<:AbstractString,1})
    # For passing list of filenames ["20150101_20160101.int", ...]
    # `collect` used to make into an array of chars for strip
    date_pairs = [split(strip(d, collect(".int")), "_")[1:2] for d in date_pairs]
    date_tups = [Tuple(d) for d in date_pairs]
    return parse_intlist_strings(date_tups)
end

parse_intlist_strings(date_pairs) = [_parse.(pair) for pair in date_pairs]

parse_geolist_strings(geolist_str::AbstractArray{String}) = _parse.(geolist_str)

import Base.string
string(arr::AbstractArray{Date}, fmt = "yyyymmdd") = Dates.format.(arr, fmt)
string(arr::AbstractArray{Tuple{Date,Date}}, fmt = "yyyymmdd") = [Dates.format.(tup, fmt) for tup in arr]

"""
    save_geolist_to_h5(h5file::String, geolist::AbstractArray{Date}; overwrite=false)

Save the geolist as a list of strings to an dataset `h5file`"""
function save_geolist_to_h5(h5file::String, geolist::AbstractArray{Date}; overwrite = false)
    !check_dset(h5file, GEOLIST_DSET, overwrite) && return
    h5write(h5file, GEOLIST_DSET, string(geolist))
end

"""In this version, save the geolist to an attribute of `object` already in `h5file`"""
function save_geolist_to_h5(h5file::String, object::String, geolist::AbstractArray{Date}; overwrite = false)
    h5open(h5file, "cw") do f
        obj = f[object]
        # Delete if exists
        # TODO: wrap this attr one into check_dset
        overwrite && exists(attrs(obj), GEOLIST_DSET) && a_delete(obj, GEOLIST_DSET)

        attrs(obj)[GEOLIST_DSET] = string(geolist)
    end
end

"""Save the geolist as a list of strings to an dataset `h5file`"""
function save_intlist_to_h5(h5file::String, intlist::AbstractArray{Tuple{Date,Date}}; overwrite = false)
    !check_dset(h5file, INTLIST_DSET, overwrite) && return
    h5write(h5file, INTLIST_DSET, _intlist_to_strings(intlist))
end
#
_intlist_to_strings(intlist::AbstractArray{Tuple{Date,Date}})::AbstractArray{String,2} = hcat([collect(tup) for tup in string(intlist)]...)

# In HDF5, DemRscs are stored as JSON dicts
load_dem_from_h5(h5file, dset = DEM_RSC_DSET) = DemRsc(JSON.parse(h5read(h5file, dset)))

function save_dem_to_h5(h5file, demrsc::DemRsc, dset_name = DEM_RSC_DSET; overwrite = true)
    !check_dset(h5file, dset_name, overwrite) && return
    h5write(h5file, dset_name, JSON.json(demrsc))
end
# TODO:
# save_dem_to_h5(h5file, demrsc::DemRsc, dset_name=DEM_RSC_DSET; overwrite=true)

function save_hdf5_stack(h5file::AbstractString, dset_name::AbstractString, stack; overwrite::Bool = false, do_permute = true)
    # TODO: is there a way to combine this with normal saving of files?
    mode = overwrite ? "w" : "cw"  # Append mode is 'cw'  (why different than "a"?)
    h5open(h5file, mode) do f
        do_permute ? write(f, dset_name, permutedims(stack, (2, 1, 3))) :
                     write(f, dset_name, stack)
    end
end

    """Wrapper around h5read to account for the transpose
necessary when reading from python-written stacks
"""
function load_hdf5_stack(h5file::AbstractString, dset_name::AbstractString; do_permute = true)
    h5open(h5file) do f
        return do_permute ? permutedims(read(f[dset_name]), (2, 1, 3)) : read(f[dset_name])
    end
end

# If loading only certain layers, don't read all into memory
function load_hdf5_stack(h5file::AbstractString, dset_name::AbstractString, valid_layer_idxs; do_permute = true)
    nrows, ncols, _ = size(h5file, dset_name)
    h5open(h5file) do f
        dset = f[dset_name]
        out = Array{eltype(dset),ndims(dset)}(undef, (nrows, ncols, length(valid_layer_idxs)))
        for (tot_idx, v_idx) in enumerate(valid_layer_idxs)
            out[:, :, tot_idx] = dset[:, :, v_idx]
        end
        return do_permute ? permutedims(out, (2, 1, 3)) : out
    end
end

"""Sum the 3rd dim (layers) of a stack without loading all into memory"""
function sum_hdf5_stack(h5file::AbstractString, dset_name::AbstractString, valid_layer_idxs)
    rows, cols, _ = size(h5file, dset_name)
    # For summing, using int instead of bool
    dtype = eltype(h5file, dset_name) <: Bool ? Int : eltype(h5file, dset_name)
    out = zeros(dtype, (rows, cols))

    h5open(h5file) do f
        d = f[dset_name]
        for ii in valid_layer_idxs
            out .+= @view d[:, :, ii][:, :, 1]
        end
    end
    return out
    end

"""Get the composite mask from the stack, true only where ALL pixels are masked"""
function load_mask(geolist::AbstractArray{Date,1}; do_permute::Bool = true, fname = "masks.h5", dset_name = "geo")
    geolist_full = load_geolist_from_h5(fname)
    idxs = indexin(geolist, geolist_full)
    out = convert(Array{Bool}, sum_hdf5_stack(fname, dset_name, idxs) .> 0)
    return do_permute ? permutedims(out) : out
end

# TODO: this sucks and doesnt work well... if one day masks out half the map,
# the geo_sum will be 0 on half, 1 on half. if we ignore that date... this still
# says "it equals the maximum of the sum, so it's masked
"""Get the composite mask from the stack, true only where ALL pixels are masked"""
function load_mask(;do_permute::Bool = true, fname = "masks.h5", dset = "geo_sum")
    mask = h5read(fname, dset)
    mval = max(1, maximum(mask))  # Make sure we dont mask all 0s
    return do_permute ? permutedims(mask .== mval) : mask .== mval
end

"""Returns false if the dataset exists and overwrite is false

If overwrite is set to true, will delete the dataset to make
sure a new one can be created"""
function check_dset(h5file, dset_name, overwrite)
    h5open(h5file, "cw") do f
        if dset_name in names(f)
            println("$dset_name already exists in $h5file")
            if overwrite
                println("Overwrite true: Deleting.")
                o_delete(f, dset_name)
            else
                println("Skipping $dset_name: not deleting")
                return false
            end
        end

        return true
    end
end

# TODO: probably a better way to do this.. but can't figure out how to
# without allocating so many arrays that it's slow as Python
function load_stack(; file_list::Union{AbstractArray{AbstractString},Nothing} = nothing,
                    directory::Union{AbstractString,Nothing} = nothing,
                    file_ext::Union{AbstractString,Nothing} = nothing)
    if isnothing(file_list)
        file_list = find_files(file_ext, directory)
    end

    # demrsc = load(find_rsc_file(basepath=directory))
    test_arr = load(file_list[1])
    rows, cols = size(test_arr)
    T = eltype(test_arr)

    stack, buffer = _return_array(T, rows, cols, length(file_list))
    for (idx, f) in enumerate(file_list)
         read!(f, buffer)
            stack[:, :, idx] = buffer
    end

    return _permute(stack, cols)
end

# Using multiple dispatch to avoid if statements for 2 types of images
function _return_array(T::Union{Type{ComplexF32},Type{Bool}}, rows::Int, cols::Int, file_len::Int)
    stack = Array{T,3}(undef, (cols, rows, file_len))
    buffer = Array{T,2}(undef, (cols, rows))
    stack, buffer
    end

function _return_array(T::Type{Float32}, rows::Int, cols::Int, file_len::Int)
    stack = Array{T,3}(undef, (2cols, rows, file_len))
    buffer = Array{T,2}(undef, (2cols, rows))
    stack, buffer
    end

# For normal .int, .geo complex/ masks, just transpose stack
function _permute(stack::AbstractArray{T,3}, cols::Int) where {T <: Number}
    return permutedims(stack, (2, 1, 3))
end

# For normal weird stacked types, pick just the right half
function _permute(stack::AbstractArray{Float32,3}, cols::Int)
    return permutedims(stack[cols + 1:end, :, :], (2, 1, 3))
end

# function _load_stack_complex(file_list::AbstractArray{AbstractString}, rows::Int, cols::Int)
#     stack = Array{ComplexF32, 3}(undef, (cols, rows, length(file_list)))
#     buffer = Array{ComplexF32, 2}(undef, (cols, rows))
#
#     for (idx, f) in enumerate(file_list)
#          read!(f, buffer)
#             stack[:, :, idx] = buffer
#     end
#
#     return permutedims(stack, (2, 1, 3))
# end
#
#
# function _load_stack_stacked(file_list::AbstractArray{AbstractString}, rows::Int, cols::Int)
#     stack = Array{Float32, 3}(undef, (2cols, rows, length(file_list)))
#     buffer = Array{Float32, 2}(undef, (2cols, rows))
#
#     for (idx, f) in enumerate(file_list)
#          read!(f, buffer)
#             stack[:, :, idx] = buffer
#     end
#     return permutedims(stack[cols+1:end, :, :], (2, 1, 3))
# end

"""
    save(filename::AbstractString, array; do_permute=true, kwargs...)

Saves a 2D/3D array as `filename`.
If the file was loaded with `do_permute=false`, it should also be
saved as `do_permute=false`

For stacked files, the expected format of `array` is a 3D stack with
`array[:,:,1]` being the amplitude image, `array[:,:,2]` being the data.
If only a 2D image is passed, an dummy amplitude image of all 1s is made.
"""
function save(filename::AbstractString, array; do_permute = true, kwargs...)
    ext = get_file_ext(filename)

    if ext == ".rsc"
        write(filename, string(array))  # array here should be a DemRsc
    elseif ext in BOOL_EXTS
        tofile(filename, array, do_permute = do_permute)
    elseif (ext in vcat(COMPLEX_EXTS, REAL_EXTS, ELEVATION_EXTS)) && (!(ext in STACKED_EXTS))
        tofile(filename, _force_float32(array), do_permute = do_permute)
    elseif ext in STACKED_EXTS
        # ndims(array) != 3 && throw(DimensionMismatch("array must be 3D [amp; data] to save as $filename"))
        if ndims(array) == 3
            amp = view(array, :, :, 1)
            data = view(array, :, :, 2)
        else
            println("!!! Warning: saving $filename with 1s for amplitude")
            data = array
            amp = ones(size(array))
        end
        # Handle the permuting for stacked here outside of function
        out = do_permute ? transpose(hcat(amp, data)) : vcat(amp, data)
        tofile(filename, _force_float32(out), do_permute = false)
    else
        error("Filetype not implemented: $filename")
    end
end

function tofile(filename, array; do_permute = true)  # , overwrite=true)
    # mode = overwrite ? "w" : "a"
    mode = "w"
    open(filename, mode) do f
        do_permute ? write(f, transpose(array)) : write(f, array)
    end
end

_iscomplex(array) = eltype(array) <: Complex
# _isfloat(array) = eltype(array) <: AbstractFloat

function _force_float32(array)
    if eltype(array) <: Complex
        return ComplexF32.(array)
    elseif eltype(array) <: AbstractFloat
        return Float32.(array)
    else
        return array
    end
end

"""Downsample a matrix by summing blocks of (row_looks, col_looks)

Cuts off values if the size isn't divisible by num looks
size = floor(rows / row_looks, cols / col_looks)
"""
function take_looks(image::AbstractArray{T}, row_looks, col_looks) where {T <: Number}
    (row_looks == 1 && col_looks == 1) && return image

    if ndims(image) > 2
        return cat((take_looks(image[:, :, i], row_looks, col_looks)
                    for i in 1:size(image, 3))..., dims = 3)
    end

    nrows, ncols = size(image)
    nr = div(nrows, row_looks)
    nc = div(ncols, col_looks)

    if T == Bool
        # Can't sum into a Bool array
        outtype = Float32
    elseif T <: Integer
        # Incase we take many looks for Int16 overflow, and to avoid InexactError
        outtype = Float64
    else
        outtype = T
    end
    out = zeros(outtype, nr, nc)
    return take_looks!(out, image, row_looks, col_looks)
end

function take_looks!(out::AbstractArray{S}, image::AbstractArray{T}, row_looks, col_looks) where {S,T <: Number}
    # Check that output is compatible
    nrows, ncols = size(image)
    nr = div(nrows, row_looks)
    nc = div(ncols, col_looks)
    (nr, nc) != size(out) && error("""
        Size of `out` $(size(out)) doesnt match taking ($row_looks, $col_looks)
        looks on image size $(size(image))""")
    return take_looks!(out, image, row_looks, col_looks, nr, nc)
end

function take_looks!(out::AbstractArray{S}, image::AbstractArray{T}, row_looks, col_looks, nr, nc) where {S,T <: Number}
    @inbounds Threads.@threads for j = 1:nc
        @inbounds for i = 1:nr
            indx_i = 1 + (i - 1) * row_looks:i * row_looks
            indx_j = 1 + (j - 1) * col_looks:j * col_looks
            @views out[i, j] = sum(image[indx_i, indx_j])
        end
    end

    out ./= (row_looks * col_looks)
    if T == Bool
        return Bool.(out .> 0)
    elseif T <: Integer
        return T.(round.(out))
    else
        return T.(out)
    end
end

function take_looks(demrsc::DemRsc, row_looks, col_looks)
    nrows, ncols = size(demrsc)
    newr = div(nrows, row_looks)
    newc = div(ncols, col_looks)
    new_x_step = demrsc.x_step * col_looks
    new_y_step = demrsc.y_step * row_looks
    return DemRsc(demrsc, file_length = newr, rows = newr, y_step = new_y_step,
                  width = newc, cols = newc, x_step = new_x_step)
end

# If we pass anything else:
# take_looks(other, row_looks, col_looks) = other

# For future:
# cc_patch(a, b) = real(abs(sum(a .* conj.(b))) / sqrt(sum(a .* conj.(a)) * sum(b .* conj.(b))))

"""
    find_igrams(;directory::AbstractString=".", parse::Bool=true, filename::AbstractString="")

Reads the list of igrams to returns Array of Igrams

Args:
    directory (str): path to the igram directory
    parse (bool): output as parsed datetime tuples. False returns the filenames
    filename (string): name of a file with .geo filenames

Returns:
    tuple(date, date) of (early, late) dates for all igrams (if parse=True)
        if parse=False: returns list[str], filenames of the igrams
"""
function find_igrams(;directory::AbstractString = ".", parse::Bool = true, filename::AbstractString = "")
    igram_file_list = isempty(filename) ? find_files("*.int", directory) : readlines(filename)
    intlist_strings = [splitpath(fname)[end] for fname in igram_file_list]
    isempty(intlist_strings) && return intlist_strings
    return parse ? parse_intlist_strings(intlist_strings) : intlist_strings
end

"""Reads in the list of .geo files used, in time order

Can also pass a filename containing .geo files as lines.

Args:
    directory (str): path to the geolist file or directory
    parse (bool): output as parsed datetime tuples. False returns the filenames
    filename (string): name of a file with .geo filenames

Returns:
    list[date]: the parse dates of each .geo used, in date order

"""
# TODO: parsers.Sentinel...
function find_geos(;directory::AbstractString = ".", parse::Bool = true, filename::AbstractString = "")
    geo_file_list = isempty(filename) ? find_files("*.geo", directory) : readlines(filename)

    !parse && return geo_file_list

    # Stripped of path for parser
    geolist_strings = [splitpath(fname)[end] for fname in geo_file_list]

    return sort(parse_geolist_strings(_strip_geoname.(geolist_strings)))

    # return sorted(parse_geolist_strings([_strip_geoname(geo) for geo in geolist_strings]))
    # if re.match(r'S1[AB]_\d{8}\.geo', geolist_strings[0]):  # S1A_YYYYmmdd.geo
    #     return sorted([_parse(_strip_geoname(geo)) for geo in geolist_strings])
    # elseif re.match(r'\d{8}', geolist_strings[0]):  # YYYYmmdd , just a date string
    #     return sorted([_parse(geo) for geo in geolist_strings])
    # else  # Full sentinel product name
    #     return sorted([apertools.parsers.Sentinel(geo).start_time.date() for geo in geolist_strings])
    # end
end

"""Leaves just date from format S1A_YYYYmmdd.geo"""
_strip_geoname(name) = reduce(replace, ["S1A_" => "", "S1B_" => "", ".geo" => ""], init = name)

end # module


using CSV, DataFrames,FreqTables, Plots, StatsPlots, StatsBase, Missings, StatsBase;

# setup filepaths
path_source = string(@__DIR__,"\\..\\source");
path_dev = string(@__DIR__,"\\..\\dev");

df = DataFrame(CSV.read(joinpath(path_dev, "19-21VAERSCOMB_clean.csv", DataFrame));

# Adrew's code starts here
# Convert SYMPTOMS strings back to an array of strings
a = fill([], size(df,1));

for rownumber in 1:size(df, 1)
    str_symptoms = df[rownumber,:SYMPTOMS]
    str_index_start = findfirst(isequal('['), str_symptoms)+2
    str_index_end = findfirst(isequal(']'), str_symptoms)-2
    a[rownumber] = split(df[rownumber,:SYMPTOMS][str_index_start:str_index_end], "\", \"")
end;

df.SYMPTOMS = a;

# Create dictionary for VAX_NAME:SYMPTOM
# Step 1: create the dict
vax_to_symptoms_dict = Dict{String, Set{String}}()
# Step 2: populate the keys (VAERS_ID) of the dict
for rownumber in 1:size(df, 1)
    vax_name = df[rownumber, :VAX_NAME]
    if !haskey(vax_to_symptoms_dict, vax_name)
        # this is the set where we will store all of the symptoms for this VAX_NAME
        vax_to_symptoms_dict[vax_name] = Set{String}()
    end
end
# Step 3: populate the values (SYMPTOMS) of the dict
for rownumber in 1:size(df, 1)
    vax_name = df[rownumber, :VAX_NAME]
    symptoms = df[rownumber, :SYMPTOMS]
    for symptom in symptoms
        push!(vax_to_symptoms_dict[vax_name], symptom)
    end
end
# View dict
vax_to_symptoms_dict

function get_freqtable(df, vax_name, symptom)
    temp = select(df, ["VAX_NAME", "SYMPTOMS"])
    # Create dummy variable for vax_name
    temp.TARGET_VAX= (temp.VAX_NAME .== vax_name)
    # Create dummy variable for symptom
    temp.TARGET_SYMPTOM = zeros(size(temp, 1))
    for rownumber in 1:size(temp, 1)
        if symptom in temp[rownumber,:SYMPTOMS]
            temp[rownumber,:TARGET_SYMPTOM] = 1
        end
    end

    # Create frequency table
    tbl = freqtable(temp, :TARGET_VAX, :TARGET_SYMPTOM)
    return tbl
end;

function get_PRR(tbl)
    # Calculate PRR
    a = tbl[2,2]
    b = tbl[2,1]
    c = tbl[1,2]
    d = tbl[1,1]
    PRR = (a/(a+b))/(c/(c+d))

    return PRR
end;

# Andrews code ends here

#=
Analysing top 5 symptoms in top 5 vaccines
COVID19 (COVID19 (PFIZER-BIONTECH)), ZOSTER (SHINGRIX), COVID19 (COVID19 (MODERNA)),
PNEUMO (PNEUMOVAX), COVID19 (COVID19 (JANSSEN))
=#

list_symptoms_mostlikely = Dict{String, Vector{String}}()

for vax in ["COVID19 (COVID19 (PFIZER-BIONTECH))",
            "ZOSTER (SHINGRIX)",
            "COVID19 (COVID19 (MODERNA))",
            "PNEUMO (PNEUMOVAX)",
            "COVID19 (COVID19 (JANSSEN))"]
    # Create dict of most reported symptoms for each vaccine
    df_vax = filter(row -> row.VAX_NAME == vax, df)

    symptoms_all_dupes = reduce(vcat, df_vax.SYMPTOMS)
    symptoms_freq_dict = StatsBase.countmap(symptoms_all_dupes)
    list_symptoms = unique(reduce(vcat, df_vax.SYMPTOMS))
    freq_per_symptom = [symptoms_freq_dict[val] for val in list_symptoms]
    perm = sortperm(freq_per_symptom; rev=true)
    # only pick the top 5 symptoms for each of the vaccine
    list_symptoms_mostlikely[vax] = list_symptoms[perm[1:5]]
end

prr_dict = Dict()

for (vax, symptoms) in list_symptoms_mostlikely
    symp_prr = []
    for symptom in symptoms
        push!(symp_prr, string(symptom, "(PRR: ",
                round(get_PRR(get_freqtable(df, vax, symptom)); digits=2), ")"))
    end
    prr_dict[vax] = symp_prr
end
prr_dict


using NiLangCore, Test

@testset "update field" begin
    @test NiLangCore.field_update(1+2im, Val(:im), 4) == 1+4im
    struct TestUpdateField1{A, B}
        a::A
    end
    @test NiLangCore.field_update(TestUpdateField1{Int,Float64}(1), Val(:a), 4) == TestUpdateField1{Int,Float64}(4)
    struct TestUpdateField2{A}
        a::A
        function TestUpdateField2(a::T) where T
            new{T}(a)
        end
    end
    @test NiLangCore.field_update(TestUpdateField2(1), Val(:a), 4) == TestUpdateField2(4)

    @test NiLangCore.default_constructor(ComplexF64, 1.0, 2.0) == 1+2im
end

@testset "_zero" begin
    @test _zero(Tuple{Float64, Float32,String,Matrix{Float64},Char,Dict{Int,Int}}) == (0.0, 0f0, "", zeros(0,0), '\0', Dict{Int,Int}())
    @test _zero(ComplexF64) == 0.0 + 0.0im
    @test _zero((1,2.0,"adsf",randn(2,2),'d',Dict(2=>5))) == (0, 0.0,"",zeros(2,2),'\0',Dict(2=>0))
    @test _zero(1+2.0im) == 0.0 + 0.0im
    @test _zero(()) == ()
    @test _zero((1,2)) == (0, 0)
    @test _zero(Symbol) == Symbol("")
    @test _zero(:x) == Symbol("")
end

@testset "fields" begin
    @test (@fields 1+3im) == (1,3)
end


module GaOperators

import MachineGa.Chromosome
import MachineGa.FloatOperators

function onepointcrossover(c1::Chromosome, c2::Chromosome)::Chromosome
    Chromosome(FloatOperators.onepointcrossover(c1.genes, c2.genes), Inf64)
end

function uniformcrossover(c1::Chromosome, c2::Chromosome)::Chromosome
    Chromosome(FloatOperators.uniformcrossover(c1.genes, c2.genes), Inf64)
end

function randombytemutation(c::Chromosome; mutateprob::Float64 = 0.10)::Chromosome
    Chromosome(FloatOperators.randombytemutation(c.genes, mutateprob = mutateprob), Inf64)
end

function nearbytemutation(c::Chromosome; mutateprob::Float64 = 0.10)::Chromosome
    Chromosome(FloatOperators.nearbytemutation(c.genes, mutateprob = mutateprob), Inf64)
end

function tournamentselection(cs::Array{Chromosome,1}; tournaments::Int = 2)::Chromosome
    csample = rand(cs, tournaments)
    bestindex = sortperm(cs, by = x -> x.cost) |> first
    best = cs[bestindex]
    return Chromosome(best.genes, best.cost)
end

end # module


#Written By Johnathan Bizzano

module HyperSphere
    import Pkg
    using Reexport
    using Lazy
    using CUDA

    const Double = Float64
    abstract type AbstractObject end

    export AbstractObject, Double

    Base.show(buffer::IO, x::AbstractObject) = print(buffer, string(x))
    Base.show(io::IO, x::BitVector) = for i in 1:length(x); print(io, x[i] ? "1" : "0"); end

    include("Utilities/Utils.jl")
    @reexport using .Utils

    include("Computation/Computation.jl")
    @reexport using .Computation

    include("Math/Math.jl")
    @reexport using .HSMath

    include("HyperDimensional/HyperDimensional.jl")
    @reexport using .HyperDimensional

    include("Data/Data.jl")
    @reexport using .Data

    include("Functions/Functions.jl")
    @reexport using .Functions

    include("NeuralNet/NeuralNet.jl")
    @reexport using .NeuralNet

    function __init__()
        __init_device__()
    end
end


using Oceananigans.Architectures: AbstractGPUArchitecture

#####
##### In place reductions!
#####

for function_name in (:sum, :prod, :maximum, :minimum, :all, :any)

    function_name! = Symbol(function_name, '!')

    # Unwrap ReducedField to a view over interior nodes:
    @eval begin
        Base.$(function_name!)(f::Function, r::AbstractReducedField, a::AbstractArray; kwargs...) = Base.$(function_name!)(f, interior(r), a; kwargs...)
        Base.$(function_name!)(r::AbstractReducedField, a::AbstractArray; kwargs...) = Base.$(function_name!)(identity, interior(r), a; kwargs...)
    end
end


# generative model interface functions

"""
    generate_s{S,A}(p::Union{POMDP{S,A},MDP{S,A}}, s::S, a::A, rng::AbstractRNG)

Return the next state given current state `s` and action taken `a`.
"""
function generate_s end

"""
    generate_o{S,A,O}(p::POMDP{S,A,O}, s::S, a::A, sp::S, rng::AbstractRNG)

Return the next observation given current state `s`, action taken `a` and next state `sp`.

Usually the observation would only depend on the next state `sp`.

    generate_o{S,A,O}(p::POMDP{S,A,O}, s::S, rng::AbstractRNG)

Return the observation from the current state. This should be used to generate initial observations.
"""
function generate_o end

"""
    generate_sr{S}(p::Union{POMDP{S},MDP{S}}, s, a, rng::AbstractRNG)

Return the next state `sp` and reward for taking action `a` in current state `s`.
"""
function generate_sr end

"""
    generate_so{S,A,O}(p::POMDP{S,A,O}, s::S, a::A, rng::AbstractRNG)

Return the next state `sp` and observation `o`.
"""
function generate_so end

"""
    generate_or{S,A,O}(p::POMDP{S,A,O}, s::S, a::A, sp::S, rng::AbstractRNG)

Return the observation `o` and reward for taking action `a` in current state `s` reaching state `sp`.
"""
function generate_or end

"""
    generate_sor{S,A,O}(p::POMDP{S,A,O}, s::S, a::A, rng::AbstractRNG)

Return the next state `sp`, observation `o` and reward for taking action `a` in current state `s`.
"""
function generate_sor end

"""
    initialstate{S}(p::Union{POMDP{S},MDP{S}}, rng::AbstractRNG)

Return the initial state for the problem `p`.

Usually the initial state is sampled from an initial state distribution.
"""
function initialstate end
@deprecate initial_state initialstate


mutable struct ForwardRelaxer{TR <: TangentRenormalizer,
                              TT <: Real,
                              } <: AbstractStage
    le_solver::TR
    t_forward_tran::TT
end

ForwardRelaxer(prob::CLVProblem, ::CLVProblem, ::CLVSolution) =
    ForwardRelaxer(prob)
ForwardRelaxer(prob::CLVProblem) = ForwardRelaxer(get_le_solver(prob),
                                                  prob.t_forward_tran)

stage_index(frx::ForwardRelaxer) = frx.le_solver.i
Base.length(frx::ForwardRelaxer) =
    ceil(Int, frx.t_forward_tran / frx.le_solver.t_renorm)
step!(frx::ForwardRelaxer) = step!(frx.le_solver)

mutable struct ForwardDynamics{S <: CLVSolution,
                               T <: TangentRenormalizer,
                               } <: AbstractStage
    le_solver::T
    i::Int
    sol::S
    R_history::Vector{UTM}

    ForwardDynamics(le_solver::T, sol::S) where {S, T} =
        new{S, T}(le_solver, 0, sol, sol.R_history)
end

ForwardDynamics(frx::ForwardRelaxer, ::CLVProblem, sol::CLVSolution) =
    ForwardDynamics(frx.le_solver, sol)

Base.length(fitr::ForwardDynamics) = length(fitr.R_history)

@inline function step!(fitr::ForwardDynamics)
    step!(fitr.le_solver)
    i = (fitr.i += 1)
    # TODO: this assignment check if lower half triangle is zero; skip that
    fitr.R_history[i] .= CLV.R_prev(fitr)
    record!(fitr, Val{:G})
    record!(fitr, Val{:x})
    return fitr
end

record!(fitr::ForwardDynamics{<:CLVSolG}, ::Type{Val{:G}}) =
    fitr.sol.G_history[fitr.i] .= CLV.G(fitr)

record!(fitr::ForwardDynamics{<:CLVSolX}, ::Type{Val{:x}}) =
    fitr.sol.x_history[fitr.i] .= phase_state(fitr)

const ForwardPass = Union{ForwardRelaxer, ForwardDynamics}

current_result(fitr::ForwardPass) = CLV.G(fitr)
phase_state(fitr::ForwardPass) = phase_state(fitr.le_solver)

mutable struct BackwardRelaxer <: AbstractStage
    num_backward_tran::Int
    R_history::Vector{UTM}
    C::UTM
    i::Int
end

BackwardRelaxer(fitr::ForwardDynamics, prob::CLVProblem, sol::CLVSolution) =
    BackwardRelaxer(sol.num_backward_tran,
                    fitr.R_history,
                    UTM(eye(fitr.R_history[end])),  # TODO: randomize
                    -1)

Base.length(brx::BackwardRelaxer) = brx.num_backward_tran

mutable struct BackwardDynamics{with_D,
                                S <: CLVSolution,
                                } <: AbstractStage
    sol::S
    R_history::Vector{UTM}
    C::UTM
    i::Int
    D_diag::Vector{Float64}

    function BackwardDynamics{with_D}(sol::S, R_history, C, i = -1,
                                      ) where {with_D, S}
        if sol isa CLVSolD && ! with_D
            error("with_D (= $with_D) has to be true when :D is",
                  " in record flag.")
        end
        bitr = new{with_D, S}(sol, R_history, C, i)
        if sol isa CLVSolC
            bitr.sol.C_history[end] .= CLV.C(bitr)
        end
        if sol isa CLVSolD
            bitr.sol.D_history[end] .= NaN
        end
        if with_D
            bitr.D_diag = similar(C, size(C, 1))
        end
        return bitr
    end
end

const BackwardDynamicsRecC = BackwardDynamics{with_D, <:CLVSolC} where {with_D}
const BackwardDynamicsRecD = BackwardDynamics{true, <:CLVSolD}
const BackwardDynamicsWithD = BackwardDynamics{true}

BackwardDynamics(brx::BackwardRelaxer, args...) =
    BackwardDynamics{false}(brx, args...)
BackwardDynamics{with_D}(brx::BackwardRelaxer,
                         ::CLVProblem,
                         sol,
                         ) where {with_D} =
    BackwardDynamics{with_D}(
        sol,
        brx.R_history[1:end - brx.i - 1],
        brx.C)

Base.length(bitr::BackwardDynamics) = length(bitr.R_history)

const BackwardPass = Union{BackwardRelaxer, BackwardDynamics}

stage_index(stage::BackwardPass) = stage.i + 1  # TODO: don't

@inline function step!(bitr::BackwardPass)
    i = (bitr.i += 1)
    R = bitr.R_history[end-i]
    C = bitr.C

    # C₀ = R⁻¹ C₁ D  (Eq. 32, Ginelli et al., 2013)
    A_ldiv_B!(R, C)
    # now:  C = R⁻¹ C₁
    for i in 1:size(C, 1)
        # C[:, i] /= norm(@view C[:, i])
        C[:, i] /= Dᵢᵢ⁻¹(bitr, i)
    end
    # now:  C = C₀ = R⁻¹ C₁ D

    record!(bitr, Val{:C})
    record!(bitr, Val{:D})
end

@inline Dᵢᵢ⁻¹(C::AbstractArray, i) = norm(@view C[:, i])
@inline Dᵢᵢ⁻¹(bitr::BackwardPass, i) = Dᵢᵢ⁻¹(bitr.C, i)
@inline function Dᵢᵢ⁻¹(bitr::BackwardDynamicsWithD, i)
    di = Dᵢᵢ⁻¹(bitr.C, i)
    bitr.D_diag[i] = 1 / di
    return di
end

function record!(bitr::BackwardDynamicsRecC, ::Type{Val{:C}})
    n = length(bitr.sol.C_history) - bitr.i - 1
    if n > 0
        bitr.sol.C_history[n] .= CLV.C(bitr)
    end
    # TODO: maybe save n=0 too
    # TODO: this assignment check if lower half triangle is zero; skip that
end

function record!(bitr::BackwardDynamicsRecD, ::Type{Val{:D}})
    n = length(bitr.sol.D_history) - bitr.i - 1
    if n > 0
        bitr.sol.D_history[n] .= bitr.D_diag
    end
    # TODO: maybe save n=0 too
end

current_result(bitr::BackwardPass) = bitr.C


#######################################################################
# Onsite/Hopping selectors
#######################################################################
abstract type Selector{M,S} end

struct OnsiteSelector{M,S} <: Selector{M,S}
    region::M
    sublats::S  # NTuple{N,NameType} (unresolved) or Vector{Int} (resolved on a lattice)
end

struct HoppingSelector{M,S,D,T} <: Selector{M,S}
    region::M
    sublats::S  # NTuple{N,Tuple{NameType,NameType}} (unres) or Vector{Tuple{Int,Int}} (res)
    dns::D
    range::T
end

"""
    onsiteselector(; region = missing, sublats = missing)

Specifies a subset of onsites energies in a given hamiltonian. Only sites at position `r` in
sublattice with name `s::NameType` will be selected if `region(r) && s in sublats` is true.
Any missing `region` or `sublat` will not be used to constraint the selection.

# See also:
    `hoppingselector`, `onsite`, `hopping`
"""
onsiteselector(; region = missing, sublats = missing) =
    OnsiteSelector(region, sanitize_sublats(sublats))

"""
    hoppingselector(; region = missing, sublats = missing, dn = missing, range = missing)

Specifies a subset of hoppings in a given hamiltonian. Only hoppings between two sites at
positions `r₁ = r - dr/2` and `r₂ = r + dr`, belonging to unit cells at integer
distance `dn´` and to sublattices `s₁` and `s₂` will be selected if: `region(r, dr) && s in
sublats && dn´ in dn && norm(dr) <= range`. If any of these is `missing` it will not be used
to constraint the selection.

# See also:
    `onsiteselector`, `onsite`, `hopping`
"""
hoppingselector(; region = missing, sublats = missing, dn = missing, range = missing) =
    HoppingSelector(region, sanitize_sublatpairs(sublats), sanitize_dn(dn), sanitize_range(range))

sanitize_sublats(s::Missing) = missing
sanitize_sublats(s::Integer) = (nametype(s),)
sanitize_sublats(s::NameType) = (s,)
sanitize_sublats(s::Tuple) = nametype.(s)
sanitize_sublats(s::Tuple{}) = ()
sanitize_sublats(n) = throw(ErrorException(
    "`sublats` for `onsite` must be either `missing`, an `s` or a tuple of `s`s, with `s::$NameType` is a sublattice name"))

sanitize_sublatpairs(s::Missing) = missing
sanitize_sublatpairs((s1, s2)::NTuple{2,Union{Integer,NameType}}) = ((nametype(s1), nametype(s2)),)
sanitize_sublatpairs((s2, s1)::Pair) = sanitize_sublatpairs((s1, s2))
sanitize_sublatpairs(s::Union{Integer,NameType}) = sanitize_sublatpairs((s,s))
sanitize_sublatpairs(s::NTuple{N,Any}) where {N} =
    ntuple(n -> first(sanitize_sublatpairs(s[n])), Val(N))
sanitize_sublatpairs(s) = throw(ErrorException(
    "`sublats` for `hopping` must be either `missing`, a tuple `(s₁, s₂)`, or a tuple of such tuples, with `sᵢ::$NameType` a sublattice name"))

sanitize_dn(dn::Missing) = missing
sanitize_dn(dn::Tuple{Vararg{NTuple{N}}}) where {N} = SVector{N,Int}.(dn)
sanitize_dn(dn::Tuple{Vararg{Number,N}}) where {N} = (SVector{N,Int}(dn),)
sanitize_dn(dn::Tuple{}) = ()

sanitize_range(::Missing) = missing
sanitize_range(range::Real) = isfinite(range) ? float(range) + sqrt(eps(float(range))) : float(range)

sublats(s::OnsiteSelector{<:Any,Missing}, lat::AbstractLattice) = collect(1:nsublats(lat))

function sublats(s::OnsiteSelector{<:Any,<:Tuple}, lat::AbstractLattice)
    names = lat.unitcell.names
    ss = Int[]
    for name in s.sublats
        i = findfirst(isequal(name), names)
        i !== nothing && push!(ss, i)
    end
    return ss
end

sublats(s::HoppingSelector{<:Any,Missing}, lat::AbstractLattice) =
    vec(collect(Iterators.product(1:nsublats(lat), 1:nsublats(lat))))

function sublats(s::HoppingSelector{<:Any,<:Tuple}, lat::AbstractLattice)
    names = lat.unitcell.names
    ss = Tuple{Int,Int}[]
    for (n1, n2) in s.sublats
        i1 = findfirst(isequal(n1), names)
        i2 = findfirst(isequal(n2), names)
        i1 !== nothing && i2 !== nothing && push!(ss, (i1, i2))
    end
    return ss
end

# selector already resolved for a lattice
sublats(s::Selector{<:Any,<:Vector}, lat) = s.sublats

# API

resolve(s::HoppingSelector, lat::AbstractLattice) =
    HoppingSelector(s.region, sublats(s, lat), _checkdims(s.dns, lat), s.range)
resolve(s::OnsiteSelector, lat::AbstractLattice) = OnsiteSelector(s.region, sublats(s, lat))

_checkdims(dns::Missing, lat::Lattice{E,L}) where {E,L} = dns
_checkdims(dns::Tuple{Vararg{SVector{L,Int}}}, lat::Lattice{E,L}) where {E,L} = dns
_checkdims(dns, lat::Lattice{E,L}) where {E,L} =
    throw(DimensionMismatch("Specified cell distance `dn` does not match lattice dimension $L"))

# are sites at (i,j) and (dni, dnj) or (dn, 0) selected?
(s::OnsiteSelector)(lat::AbstractLattice, (i, j)::Tuple, (dni, dnj)::Tuple{SVector, SVector}) =
    isonsite((i, j), (dni, dnj)) && isinregion(i, dni, s.region, lat) && isinsublats(sublat(lat, i), s.sublats)

(s::HoppingSelector)(lat::AbstractLattice, inds, dns) =
    !isonsite(inds, dns) && isinregion(inds, dns, s.region, lat) && isindns(dns, s.dns) &&
    isinrange(inds, s.range, lat) && isinsublats(sublat.(Ref(lat), inds), s.sublats)

isonsite((i, j), (dni, dnj)) = i == j && dni == dnj

isinregion(i::Int, dn, ::Missing, lat) = true
isinregion(i::Int, dn, region::Function, lat) = region(sites(lat)[i] + bravais(lat) * dn)

isinregion(is::Tuple{Int,Int}, dns, ::Missing, lat) = true
function isinregion((row, col)::Tuple{Int,Int}, (dnrow, dncol), region::Function, lat)
    br = bravais(lat)
    r, dr = _rdr(sites(lat)[col] + br * dncol, sites(lat)[row] + br * dnrow)
    return region(r, dr)
end

isinsublats(s::Int, ::Missing) = true
isinsublats(s::Int, sublats::Vector{Int}) = s in sublats
isinsublats(ss::Tuple{Int,Int}, ::Missing) = true
isinsublats(ss::Tuple{Int,Int}, sublats::Vector{Tuple{Int,Int}}) = ss in sublats
isinsublats(s, sublats) =
    throw(ArgumentError("Sublattices $sublats in selector are not resolved."))

isindns((dnrow, dncol)::Tuple{SVector,SVector}, dns) = isindns(dnrow - dncol, dns)
isindns(dn::SVector{L,Int}, dns::Tuple{Vararg{SVector{L,Int}}}) where {L} = dn in dns
isindns(dn::SVector, dns::Missing) = true
isindns(dn, dns) =
    throw(ArgumentError("Cell distance dn in selector is incompatible with Lattice."))

isinrange(inds, ::Missing, lat) = true
isinrange((row, col)::Tuple{Int,Int}, range::Number, lat) =
    norm(sites(lat)[col] - sites(lat)[row]) <= range

#######################################################################
# TightbindingModelTerm
#######################################################################
abstract type TightbindingModelTerm end
abstract type AbstractOnsiteTerm <: TightbindingModelTerm end
abstract type AbstractHoppingTerm <: TightbindingModelTerm end

struct OnsiteTerm{F,S<:OnsiteSelector,C} <: AbstractOnsiteTerm
    o::F
    selector::S
    coefficient::C
    forcehermitian::Bool
end

struct HoppingTerm{F,S<:HoppingSelector,C} <: AbstractHoppingTerm
    t::F
    selector::S
    coefficient::C
    forcehermitian::Bool
end

(o::OnsiteTerm{<:Function})(r,dr) = o.coefficient * o.o(r)
(o::OnsiteTerm)(r,dr) = o.coefficient * o.o

(h::HoppingTerm{<:Function})(r, dr) = h.coefficient * h.t(r, dr)
(h::HoppingTerm)(r, dr) = h.coefficient * h.t

sublats(t::TightbindingModelTerm, lat) = sublats(t.selector, lat)

displayparameter(::Type{<:Function}) = "Function"
displayparameter(::Type{T}) where {T} = "$T"

function Base.show(io::IO, o::OnsiteTerm{F}) where {F}
    i = get(io, :indent, "")
    print(io,
"$(i)OnsiteTerm{$(displayparameter(F))}:
$(i)  Sublattices      : $(o.selector.sublats === missing ? "any" : o.selector.sublats)
$(i)  Force hermitian  : $(o.forcehermitian)
$(i)  Coefficient      : $(o.coefficient)")
end

function Base.show(io::IO, h::HoppingTerm{F}) where {F}
    i = get(io, :indent, "")
    print(io,
"$(i)HoppingTerm{$(displayparameter(F))}:
$(i)  Sublattice pairs : $(h.selector.sublats === missing ? "any" : (t -> Pair(reverse(t)...)).(h.selector.sublats))
$(i)  dn cell distance : $(h.selector.dns === missing ? "any" : h.selector.dns)
$(i)  Hopping range    : $(round(h.selector.range, digits = 6))
$(i)  Force hermitian  : $(h.forcehermitian)
$(i)  Coefficient      : $(h.coefficient)")
end

# External API #
"""
    onsite(o; forcehermitian = true, kw...)
    onsite(o, onsiteselector(; kw...); forcehermitian = true)

Create an `TightbindingModelTerm` that applies an onsite energy `o` to a `Lattice` when
creating a `Hamiltonian` with `hamiltonian`. A subset of sites can be specified with the
`kw...`, see `onsiteselector` for details.

The onsite energy `o` can be a number, a matrix (preferably `SMatrix`) or a function of the
form `r -> ...` for a position-dependent onsite energy. If `forcehermitian` is true, the
model will produce an hermitian Hamiltonian.

The dimension of `o::AbstractMatrix` must match the orbital dimension of applicable
sublattices (see also `orbitals` option for `hamiltonian`). If `o::Number` it will be
treated as `o * I` (proportional to identity matrix) when applied to multiorbital
sublattices.

`TightbindingModelTerm`s created with `onsite` or `hopping` can be added or substracted
together to build more complicated `TightbindingModel`s.

# Examples
```
julia> onsite(1, sublats = (:A,:B)) - hopping(2, sublats = :A=>:A)
TightbindingModel{2}: model with 2 terms
  OnsiteTerm{Int64}:
    Sublattices      : (:A, :B)
    Force hermitian  : true
    Coefficient      : 1
  HoppingTerm{Int64}:
    Sublattice pairs : (:A => :A,)
    dn cell distance : any
    Hopping range    : 1.0
    Force hermitian  : true
    Coefficient      : -1

julia> LatticePresets.honeycomb() |> hamiltonian(onsite(r->@SMatrix[1 2; 3 4]), orbitals = Val(2))
Hamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 1 (SparseMatrixCSC, sparse)
  Harmonic size    : 2 × 2
  Orbitals         : ((:a, :a), (:a, :a))
  Element type     : 2 × 2 blocks (Complex{Float64})
  Onsites          : 2
  Hoppings         : 0
  Coordination     : 0.0
```

# See also:
    `hopping`, `onsiteselector`, `hoppingselector`
"""
onsite(o; forcehermitian = true, kw...) =
    onsite(o, onsiteselector(; kw...); forcehermitian = forcehermitian)
onsite(o, selector; forcehermitian::Bool = true) =
    TightbindingModel(OnsiteTerm(o, selector, 1, forcehermitian))

"""
    hopping(t; forcehermitian = true, range = 1, kw...)
    hopping(t, hoppingselector(; range = 1, kw...); forcehermitian = true)

Create an `TightbindingModelTerm` that applies a hopping `t` to a `Lattice` when
creating a `Hamiltonian` with `hamiltonian`. A subset of hoppings can be specified with the
`kw...`, see `hoppingselector` for details. Note that a default `range = 1` is assumed.

The hopping amplitude `t` can be a number, a matrix (preferably `SMatrix`) or a function
of the form `(r, dr) -> ...` for a position-dependent hopping (`r` is the bond center,
and `dr` the bond vector). If `sublats` is specified as a sublattice name pair, or tuple
thereof, `hopping` is only applied between sublattices with said names. If `forcehermitian`
is true, the model will produce an hermitian Hamiltonian.

The dimension of `t::AbstractMatrix` must match the orbital dimension of applicable
sublattices (see also `orbitals` option for `hamiltonian`). If `t::Number` it will be
treated as `t * I` (proportional to identity matrix) when applied to multiorbital
sublattices.

`TightbindingModelTerm`s created with `onsite` or `hopping` can be added or substracted
together to build more complicated `TightbindingModel`s.

# Examples
```
julia> onsite(1) - hopping(2, dn = ((1,2), (0,0)), sublats = :A=>:B)
TightbindingModel{2}: model with 2 terms
  OnsiteTerm{Int64}:
    Sublattices      : any
    Force hermitian  : true
    Coefficient      : 1
  HoppingTerm{Int64}:
    Sublattice pairs : (:A => :B,)
    dn cell distance : ([1, 2], [0, 0])
    Hopping range    : 1.0
    Force hermitian  : true
    Coefficient      : -1

julia> LatticePresets.honeycomb() |> hamiltonian(hopping((r,dr) -> cos(r[1]), sublats = ((:A,:A), (:B,:B))))
Hamiltonian{<:Lattice} : Hamiltonian on a 2D Lattice in 2D space
  Bloch harmonics  : 7 (SparseMatrixCSC, sparse)
  Harmonic size    : 2 × 2
  Orbitals         : ((:a,), (:a,))
  Element type     : scalar (Complex{Float64})
  Onsites          : 0
  Hoppings         : 12
  Coordination     : 6.0
```

# See also:
    `onsite`, `onsiteselector`, `hoppingselector`
"""
hopping(t; forcehermitian = true, range = 1, kw...) =
    hopping(t, hoppingselector(; range = range, kw...); forcehermitian = forcehermitian)
hopping(t, selector; forcehermitian = true) =
    TightbindingModel(HoppingTerm(t, selector, 1, forcehermitian))

Base.:*(x, o::OnsiteTerm) =
    OnsiteTerm(o.o, o.selector, x * o.coefficient, o.forcehermitian)
Base.:*(x, t::HoppingTerm) =
    HoppingTerm(t.t, t.selector, x * t.coefficient, t.forcehermitian)
Base.:*(t::TightbindingModelTerm, x) = x * t
Base.:-(t::TightbindingModelTerm) = (-1) * t

LinearAlgebra.ishermitian(t::OnsiteTerm) = t.forcehermitian
LinearAlgebra.ishermitian(t::HoppingTerm) = t.forcehermitian

#######################################################################
# TightbindingModel
#######################################################################
struct TightbindingModel{N,T<:Tuple{Vararg{TightbindingModelTerm,N}}}
    terms::T
end

terms(t::TightbindingModel) = t.terms

TightbindingModel(ts::TightbindingModelTerm...) = TightbindingModel(ts)

(m::TightbindingModel)(r, dr) = sum(t -> t(r, dr), m.terms)

# External API #

Base.:*(x, m::TightbindingModel) = TightbindingModel(x .* m.terms)
Base.:*(m::TightbindingModel, x) = x * m
Base.:-(m::TightbindingModel) = TightbindingModel((-1) .* m.terms)

Base.:+(m::TightbindingModel, t::TightbindingModel) = TightbindingModel((m.terms..., t.terms...))
Base.:-(m::TightbindingModel, t::TightbindingModel) = m + (-t)

function Base.show(io::IO, m::TightbindingModel{N}) where {N}
    ioindent = IOContext(io, :indent => "  ")
    print(io, "TightbindingModel{$N}: model with $N terms", "\n")
    foreach(t -> print(ioindent, t, "\n"), m.terms)
end

LinearAlgebra.ishermitian(m::TightbindingModel) = all(t -> ishermitian(t), m.terms)

#######################################################################
# offdiagonal
#######################################################################
"""
    offdiagonal(model, lat, nsublats::NTuple{N,Int})

Build a restricted version of `model` that applies only to off-diagonal blocks formed by
sublattice groups of size `nsublats`.
"""
offdiagonal(m::TightbindingModel, lat, nsublats) =
    TightbindingModel(offdiagonal.(m.terms, Ref(lat), Ref(nsublats)))

offdiagonal(o::OnsiteTerm, lat, nsublats) =
    throw(ArgumentError("No onsite terms allowed in off-diagonal coupling"))

function offdiagonal(t::HoppingTerm, lat, nsublats)
    selector´ = resolve(t.selector, lat)
    s = selector´.sublats
    sr = sublatranges(nsublats...)
    filter!(spair ->  findblock(first(spair), sr) != findblock(last(spair), sr), s)
    return HoppingTerm(t.t, selector´, t.coefficient, t.forcehermitian)
end

sublatranges(i::Int, is::Int...) = _sublatranges((1:i,), is...)
_sublatranges(rs::Tuple, i::Int, is...) = _sublatranges((rs..., last(last(rs)) + 1: last(last(rs)) + i), is...)
_sublatranges(rs::Tuple) = rs

findblock(s, sr) = findfirst(r -> s in r, sr)

#######################################################################
# checkmodelorbs - check for inconsistent orbital dimensions
#######################################################################
checkmodelorbs(model::TightbindingModel, orbs, lat) =
    foreach(term -> _checkmodelorbs(term, orbs, lat), model.terms)

function _checkmodelorbs(term::HoppingTerm, orbs, lat)
    for (s1, s2) in sublats(term, lat)
        _checkmodelorbs(term(first(sites(lat, s1)), first(sites(lat, s2))), length(orbs[s1]), length(orbs[s2]))
    end
    return nothing
end

function _checkmodelorbs(term::OnsiteTerm, orbs, lat)
    for s in sublats(term, lat)
        _checkmodelorbs(term(first(sites(lat, s)), first(sites(lat, s))), length(orbs[s]))
    end
    return nothing
end

_checkmodelorbs(s::SMatrix, m, n = m) =
    size(s) == (m, n) || @warn("Possible dimension mismatch between model and Hamiltonian. Did you correctly specify the `orbitals` in hamiltonian?")

_checkmodelorbs(s::Number, m, n = m) = _checkmodelorbs(SMatrix{1,1}(s), m, n)

#######################################################################
# onsite! and hopping!
#######################################################################
abstract type ElementModifier{V,F,S} end

struct Onsite!{V<:Val,F<:Function,S<:Selector} <: ElementModifier{V,F,S}
    f::F
    needspositions::V    # Val{false} for f(o; kw...), Val{true} for f(o, r; kw...) or other
    selector::S
    forcehermitian::Bool
    addconjugate::Bool   # determines whether to return f(o) or (f(o) + f(o')')/2
end

Onsite!(f, selector, forcehermitian = true) =
    Onsite!(f, Val(!applicable(f, 0.0)), selector, forcehermitian, false)

struct Hopping!{V<:Val,F<:Function,S<:Selector} <: ElementModifier{V,F,S}
    f::F
    needspositions::V    # Val{false} for f(h; kw...), Val{true} for f(h, r, dr; kw...) or other
    selector::S
    forcehermitian::Bool
    addconjugate::Bool   # determines whether to return f(t) or (f(t) + f(t')')/2
end

Hopping!(f, selector, forcehermitian = true) =
    Hopping!(f, Val(!applicable(f, 0.0)), selector, forcehermitian, false)

# API #

"""
    onsite!(f; forcehermitian = true, kw...)
    onsite!(f, onsiteselector(; kw...); forcehermitian = true)

Create an `ElementModifier`, to be used with `parametric`, that applies `f` to onsite
energies specified by `onsiteselector(; kw...)`. The form of `f` may be `f = (o; kw...) ->
...` or `f = (o, r; kw...) -> ...` if the modification is position (`r`) dependent. The
former is naturally more efficient, as there is no need to compute the positions of each
onsite energy. If `forcehermitian == true` the modification is forced to yield a hermitian
Hamiltonian if the original was hermitian.

# See also:
    `hopping!`, `parametric`
"""
onsite!(f; forcehermitian = true, kw...) =
    onsite!(f, onsiteselector(; kw...); forcehermitian = forcehermitian)
onsite!(f, selector; forcehermitian = true) = Onsite!(f, selector, forcehermitian)

"""
    hopping!(f; forcehermitian = true, kw...)
    hopping!(f, hoppingselector(; kw...); forcehermitian = true)

Create an `ElementModifier`, to be used with `parametric`, that applies `f` to hoppings
specified by `hoppingselector(; kw...)`. The form of `f` may be `f = (t; kw...) -> ...` or
`f = (t, r, dr; kw...) -> ...` if the modification is position (`r, dr`) dependent. The
former is naturally more efficient, as there is no need to compute the positions of the two
sites involved in each hopping. If `forcehermitian == true` the modification is forced to
yield a hermitian Hamiltonian if the original was hermitian.

# See also:
    `onsite!`, `parametric`
"""
hopping!(f; forcehermitian = true, kw...) = hopping!(f, hoppingselector(; kw...); forcehermitian = forcehermitian)
hopping!(f, selector; forcehermitian = true) = Hopping!(f, selector, forcehermitian)

function resolve(o::Onsite!, lat)
    addconjugate = o.forcehermitian
    Onsite!(o.f, o.needspositions, resolve(o.selector, lat), o.forcehermitian, addconjugate)
end

function resolve(h::Hopping!, lat)
    addconjugate = h.selector.sublats === missing && h.forcehermitian
    Hopping!(h.f, h.needspositions, resolve(h.selector, lat), h.forcehermitian, addconjugate)
end

# Intended for resolved ElementModifiers only
@inline (o!::Onsite!{Val{false}})(o, r; kw...) = o!(o; kw...)
@inline (o!::Onsite!{Val{false}})(o, r, dr; kw...) = o!(o; kw...)
@inline (o!::Onsite!{Val{false}})(o; kw...) =
    o!.addconjugate ? 0.5 * (o!.f(o; kw...) + o!.f(o'; kw...)') : o!.f(o; kw...)
@inline (o!::Onsite!{Val{true}})(o, r, dr; kw...) = o!(o, r; kw...)
@inline (o!::Onsite!{Val{true}})(o, r; kw...) =
    o!.addconjugate ? 0.5 * (o!.f(o, r; kw...) + o!.f(o', r; kw...)') : o!.f(o, r; kw...)

@inline (h!::Hopping!{Val{false}})(t, r, dr; kw...) = h!(t; kw...)
@inline (h!::Hopping!{Val{false}})(t; kw...) =
    h!.addconjugate ? 0.5 * (h!.f(t; kw...) + h!.f(t'; kw...)') : h!.f(t; kw...)
@inline (h!::Hopping!{Val{true}})(t, r, dr; kw...) =
    h!.addconjugate ? 0.5 * (h!.f(t, r, dr; kw...) + h!.f(t', r, -dr; kw...)') : h!.f(t, r, dr; kw...)


module SPN

include("nodes.jl")
export Node, InnerNode, LeafNode, SumNode, ProdNode, IndicatorNode, GaussianNode
export connect!, setInput!, eval!, passDerivative!

include("network.jl")
export SumProductNetwork, computeDerivatives!
export numNodes, numSumNodes, numProdNodes, numLeafNodes

include("inference.jl")
export marginalInference!, conditionalInference!, mpeInference!

include("utils.jl")
export addExp, quantileMeans, dataLikelihood!

include("learn/parameter.jl")
export parameterLearnHardEM!, parameterLearnEM!, parameterLearnGD!

include("learn/poon.jl")
export structureLearnPoon

end # module


samplingparameter(samplingstyle::SamplingStyle, ap::ApproximationProblem; samplingstrategy=DefaultSamplingStrategy(), options...)=
    samplingparameter(samplingstyle, samplingstrategy, ap; options...)

samplingparameter(samplingstyle::GenericSamplingStyle, ap::ApproximationProblem; options...) = -1

function samplingparameter(samplingstyle::SamplingStyle, samplingstrategy::SamplingStrategy, ap::ApproximationProblem;
        verbose=false, forget=false, options...)
    if haskey(options, :L)
        # The user specified L as an option
        L = options[:L]
        setsamplingparam!(ap, L)
    elseif samplingparam(ap) != nothing && !forget
        # Has it been computed before
        L = samplingparam(ap)
    else
        # It wasn't. We deduce its value from the options given and store the outcome.
        L = deduce_samplingparameter(samplingstyle, samplingstrategy, ap; verbose=verbose, options...)
        setsamplingparam!(ap, L)
    end
    verbose && println("Sampling parameter: using L = $L")
    return L
end

samplingparameter(samplingstyle::DiscreteGramStyle, samplingstrategy::SamplingStrategy, ap::ApproximationProblem; options...) =
    samplingparameter(OversamplingStyle(), samplingstrategy, ap; options...)

deduce_samplingparameter(ss::SamplingStyle, samplingstrategy::SamplingStrategy, ap::ApproximationProblem; options...) =
    deduce_samplingparameter(ss, samplingstrategy, platform(ap), parameter(ap); options...)

deduce_samplingparameter(ss::ProductSamplingStyle, samplingstrategy::SamplingStrategy, ap::ProductPlatformApproximation; options...) =
    map((x,style)->samplingparameter(style, samplingstrategy, x; options...), ApproximationProblems.unsafe_elements(ap), ss.styles)

function deduce_samplingparameter(ss::ProductSamplingStyle{NTuple{N,OversamplingStyle}}, samplingstrategy::SamplingStrategy, ap::ProductPlatformApproximation;
        dict=dictionary(ap), oversamplingfactor=2, verbose=false, options...) where N
    oversamplingfactor = oversamplingfactor_float2tuple(samplingstrategy, oversamplingfactor, dict)
    if haskey(options, :M)
        M = options[:M]
        if M isa Integer
            M = M_int2tuple(samplingstrategy, M, dict)
        end
    else
        M = oversamplingM(samplingstrategy, oversamplingfactor, dict)
    end
    L = tuple(map((api, Mi)->deduce_samplingparameter(OversamplingStyle(), samplingstrategy, api;
        oversamplingfactor=oversamplingfactor, M=Mi, verbose=verbose, options...), ApproximationProblems.unsafe_elements(ap), M)...)
    verbose && println("Sampling parameter: best match for M = $M is L = $L")
    L
end

M_int2tuple(::DefaultSamplingStrategy, M::Int, dict::TensorProductDict) =
    ntuple(k->round(Int,M*size(dict, k)/length(dict)))
oversamplingM(::DefaultSamplingStrategy, oversamplingfactor::Real, dict::TensorProductDict) =
    ntuple(k->round(Int,oversamplingfactor*size(dict, k)),Val(dimension(dict)))
oversamplingfactor_float2tuple(::DefaultSamplingStrategy, oversamplingfactor::Real, dict::TensorProductDict) =
    oversamplingfactor^(1/dimension(dict))

function deduce_oversampling_parameter(ss::OversamplingStyle, samplingstrategy::SamplingStrategy, args...; dict=dictionary(args...), verbose = false, oversamplingfactor = 2, options...)
    # In the absence of L, we deduce M and then find the best matching L
    # M is either supplied, or we compute it based on the (default) oversamplingfactor
    M = haskey(options, :M) ? options[:M] : round(Int, oversamplingfactor * length(dict))
    verbose && println("Sampling parameter: oversamplingfactor=$oversamplingfactor, length=$(length(dict)), M=$M")
    L = match_and_correct_sampling_parameter(samplingstrategy, args..., M; samplingstyle=ss, options...)
    verbose && println("Sampling parameter: best match for M = $M is L = $L")
    L
end

deduce_samplingparameter(ss::SamplingStyle, samplingstrategy::SamplingStrategy, platform::Platform, param; dict = dictionary(platform, param), options...) =
    deduce_samplingparameter(ss, samplingstrategy, dict; options...)
deduce_samplingparameter(ss::OversamplingStyle, samplingstrategy::SamplingStrategy, dict::Dictionary; options...) =
    deduce_oversampling_parameter(ss, samplingstrategy, dict; dict=dict, options...)
deduce_samplingparameter(ss::OversamplingStyle, samplingstrategy::SamplingStrategy, platform::Platform, param; dict=dictionary(platform, param), options...) =
    deduce_oversampling_parameter(ss, samplingstrategy, platform, param; dict=dict, options...)
deduce_samplingparameter(ss::InterpolationStyle, samplingstrategy::SamplingStrategy, dict::Dictionary; options...) =
    match_and_correct_sampling_parameter(samplingstrategy, dict, length(dict); samplingstyle=ss, options...)
# TODO: implement this one better (more general)
deduce_samplingparameter(::GramStyle, samplingstrategy::SamplingStrategy, dict::Dictionary; options...) = length(dict)

deduce_samplingparameter(::RectangularGramStyle, samplingstrategy::SamplingStrategy, dict::Dictionary; projectiondict, options...) = length(projectiondict)

deduce_samplingparameter(::GridStyle, samplingstrategy::SamplingStrategy, dict::Dictionary; options...) = param(dict)


function build_robot!(vis, model::Box;
    i=1,
    r=0.1,
    tl=1.0,
    box_color=Colors.RGBA(0.0, 0.0, 0.0, tl))

    setobject!(vis["box_$i"], GeometryBasics.Rect(
        Vec(-1.0 * r, -0.005, -1.0 * r),
        Vec(2.0 * r, 0.01, 2.0 * r)),
        MeshPhongMaterial(color=box_color))
end

function set_robot!(vis, model::Box, q; i=1)
    settransform!(vis["box_$i"],
        compose(Translation(q[1], 0.001 * i, q[2]), LinearMap(RotY(-1.0 * q[3]))))
end

function visualize!(vis, model::Box, q;
    i=1,
    r=0.1,
    tl=1.0,
    box_color=Colors.RGBA(0.0, 0.0, 0.0, tl),
    Δt=0.1,
    fixed_camera=true)

    default_background!(vis)

    build_robot!(vis, model,
        i=i,
        r=r,
        tl=tl,
        box_color= box_color)

    anim = MeshCat.Animation(convert(Int, floor(1.0 / Δt)))

    T = length(q)
    for t = 1:T-1
        MeshCat.atframe(anim, t) do
            set_robot!(vis, model, q[t], i=i)
        end
    end

    if fixed_camera
        MeshCat.settransform!(vis["/Cameras/default"],
        MeshCat.compose(MeshCat.Translation(0.0, -50.0, -1.0), MeshCat.LinearMap(Rotations.RotZ(-pi / 2.0))))
        setprop!(vis["/Cameras/default/rotated/<object>"], "zoom", 25)
    end

    MeshCat.setanimation!(vis, anim)
end


using Utils
using Underscores

nums(s) = @_ split(s) |> map(parse(Int, _), __)

struct Rule
    refs::Array{Int}
end

struct Alternative
    left::Array{Int}
    right::Array{Int}
end

const RulePossibilities = Union{AbstractString,Rule,Alternative}
const Rules = Dict{Int,RulePossibilities}

struct Messages
    rules::Rules
    messages::Array{SubString}
end

function parse_file(f)
    rules, messages = split_blocks(f)

    rls = Rules()
    for line in eachline(IOBuffer(rules))
        head, content = split(line, ": ")
        head = parse(Int, head)
        if startswith(content, '"')
            rls[head] = strip(content, '"')
        else
            if occursin('|', content)
                a, b = split(content, " | ")
                rls[head] = Alternative(nums(a), nums(b))
            else
                rls[head] = Rule(nums(content))
            end
        end
    end
    Messages(rls, collect(eachline(IOBuffer(messages))))
end

function load()
    open(joinpath(@__DIR__, "19.txt"), "r") do f
        parse_file(f)
    end
end

messages = load()

function build_regex(rules::Rules, current)
    current = rules[current]
    if isa(current, Alternative)
        lefts = @_ current.left |> map(build_regex(rules, _), __) |> join
        rights = @_ current.right |> map(build_regex(rules, _), __) |> join
        "(?:$lefts|$rights)"
    elseif isa(current, Rule)
        @_ current.refs |> map(build_regex(rules, _), __) |> join
    else
        current
    end
end

build_regex(rules) = '^' * build_regex(rules, 0) * '$' |> Regex

function ex1(messages)
    re = build_regex(messages.rules)
    @_ count(occursin(re, _), messages.messages)
end

function rule_walker(_, str, current::AbstractString, pos)::Array{Int,1}
    [startswith(str[pos:end], current) ? pos + length(current) : 0]
end

function rule_walker(rules::Rules, str, current::Int, pos)::Array{Int,1}
    rule_walker(rules, str, rules[current], pos)
end

function rule_walker(rules::Rules, str, current::Alternative, pos)::Array{Int,1}
    re = rule_walker(rules, str, Rule(current.left), pos)
    append!(re, rule_walker(rules, str, Rule(current.right), pos))
    filter(p -> p != 0, re)
end

function rule_walker(rules::Rules, str, current::Rule, pos)::Array{Int,1}
    positions = [pos]
    for (i,ref) in enumerate(current.refs)
        positions = @_ map(rule_walker(rules, str, ref, _), positions) |> Iterators.flatten |> collect
        if i < lastindex(current.refs)
            filter!(p -> p in eachindex(str), positions)
        end
    end
    positions
end

rule_walker(rules, str) = (length(str) + 1) in rule_walker(rules, str, 0, 1)

function ex2(messages, update=true)
    changed = copy(messages.rules)
    if update
        changed[8] = Alternative([42], [42, 8])
        changed[11] = Alternative([42, 31], [42, 11, 31])
    end

    @_ count(rule_walker(changed, _), messages.messages)
end

println("Number of matching messages: ", ex1(messages))
println("Matching after change: ", ex2(messages))

using Test

@testset "Monster Messages" begin
    input1 = """
    0: 4 1 5
    1: 2 3 | 3 2
    2: 4 4 | 5 5
    3: 4 5 | 5 4
    4: "a"
    5: "b"

    ababbb
    bababa
    abbbab
    aaabbb
    aaaabbb
    """

    example1 = parse_file(IOBuffer(input1))

    input2 = """
    42: 9 14 | 10 1
    9: 14 27 | 1 26
    10: 23 14 | 28 1
    1: "a"
    11: 42 31
    5: 1 14 | 15 1
    19: 14 1 | 14 14
    12: 24 14 | 19 1
    16: 15 1 | 14 14
    31: 14 17 | 1 13
    6: 14 14 | 1 14
    2: 1 24 | 14 4
    0: 8 11
    13: 14 3 | 1 12
    15: 1 | 14
    17: 14 2 | 1 7
    23: 25 1 | 22 14
    28: 16 1
    4: 1 1
    20: 14 14 | 1 15
    3: 5 14 | 16 1
    27: 1 6 | 14 18
    14: "b"
    21: 14 1 | 1 14
    25: 1 1 | 1 14
    22: 14 14
    8: 42
    26: 14 22 | 1 20
    18: 15 15
    7: 14 5 | 1 21
    24: 14 1

    abbbbbabbbaaaababbaabbbbabababbbabbbbbbabaaaa
    bbabbbbaabaabba
    babbbbaabbbbbabbbbbbaabaaabaaa
    aaabbbbbbaaaabaababaabababbabaaabbababababaaa
    bbbbbbbaaaabbbbaaabbabaaa
    bbbababbbbaaaaaaaabbababaaababaabab
    ababaaaaaabaaab
    ababaaaaabbbaba
    baabbaaaabbaaaababbaababb
    abbbbabbbbaaaababbbbbbaaaababb
    aaaaabbaabaaaaababaa
    aaaabbaaaabbaaa
    aaaabbaabbaaaaaaabbbabbbaaabbaabaaa
    babaaabbbaaabaababbaabababaaab
    aabbbbbaabbbaaaaaabbbbbababaaaaabbaaabba
    """

    example2 = parse_file(IOBuffer(input2))

    @testset "interna" begin
        @testset "string rule" begin
            @test rule_walker(Rules(), "abcd", "ab", 1) == [3]
            @test rule_walker(Rules(), "abcd", "abc", 1) == [4]
            @test rule_walker(Rules(), "abcd", "d", 4) == [5]
            @test rule_walker(Rules(), "abcd", "bc", 2) == [4]
            @test rule_walker(Rules(), "abcd", "bc", 1) == [0]
        end

        @testset "int rule" begin
            @assert rule_walker(Rules(0 => "ab"), "abcd", 0, 1) == [3]
            @assert rule_walker(Rules(1 => "abc"), "abcd", 1, 1) == [4]
            @assert rule_walker(Rules(2 => "d"), "abcd", 2, 4) == [5]
            @assert rule_walker(Rules(3 => "bc"), "abcd", 3, 2) == [4]
            @assert rule_walker(Rules(4 => "bc"), "abcd", 4, 1) == [0]
        end

        @testset "alternative rule" begin
            @assert 5 in rule_walker(Rules(0 => Alternative([1, 2], [2, 1]), 1 => "ab", 2 => "cd"), "abcd", 0, 1)
            @assert rule_walker(Rules(0 => Alternative([1, 2], [2, 1]), 1 => "ab", 2 => "cde"), "abcd", 0, 1) == []
            @assert 5 in rule_walker(Rules(0 => Alternative([1, 2], [2, 3]), 1 => "ab", 2 => "cd", 3 => "cde"), "abcd", 0, 1)
            @assert 6 in rule_walker(Rules(0 => Alternative([1, 2], [1, 3]), 1 => "ab", 2 => "cd", 3 => "cde"), "abcde", 0, 1)
            @assert 7 in rule_walker(Rules(0 => Rule([9, 10]), 9 => Alternative([1, 2], [1, 3]), 1 => "ab", 2 => "cd", 3 => "cde", 10 => "z"), "abcdez", 0, 1)
        end

        @testset "rule" begin
            @assert !(4 in rule_walker(Rules(0 => Rule([1, 2]), 1 => "a", 2 => "b"), "abc", 0, 1))
            @assert rule_walker(Rules(0 => Rule([1, 2]), 1 => "b", 2 => "a"), "ab", 0, 1) == []
            @assert 3 in rule_walker(Rules(0 => Rule([1, 2]), 1 => "a", 2 => "b"), "ab", 0, 1)
            @assert 4 in rule_walker(Rules(0 => Rule([1, 2]), 1 => "a", 2 => Alternative([3], [3, 4]), 3 => "b", 4 => "c"), "abc", 0, 1)
        end

        @testset "shell" begin
            @assert !rule_walker(Rules(0 => Rule([1, 2]), 1 => "b", 2 => "a"), "ab")
            @assert rule_walker(Rules(0 => Rule([1, 2]), 1 => "a", 2 => "b"), "ab")
            @assert !rule_walker(Rules(0 => Rule([1, 2]), 1 => "a", 2 => "b"), "abc")
            @assert rule_walker(Rules(0 => Rule([1, 2]), 1 => "a", 2 => Alternative([3], [3, 4]), 3 => "b", 4 => "c"), "abc")
        end
    end

    @testset "example 1" begin
        @test ex1(example1) == 2
        @test ex1(example2) == 3
    end

    @testset "example 2" begin
        example3 = Messages(copy(example2.rules), example2.messages)
        example3.rules[8] = Alternative([42], [42, 8])
        example3.rules[11] = Alternative([42, 31], [42, 11, 31])

        @test ex2(example1, false) == 2
        @test ex2(example2, false) == 3
        @test ex2(example3, false) == 12
    end

    @testset "results" begin
        @test ex1(messages) == 205
        @test ex2(messages) == 329
    end
end