text:
##################################################
## SymbolicObject types have field x::PyCall.PyObject

## Symbol class for controlling dispatch
abstract type SymbolicObject <: Number end
struct Sym <: SymbolicObject
    __pyobject__::PyCall.PyObject
end

## Matrices
## We use this class for `ImmutableMatrices`
## Mutable matrices are mapped to `AbstractArray{Sym,N}`
## cf. matrix.jl
"""
    SymMatrix

Type to store a SymPy matrix, as created by `sympy.ImmutableMatrix`.

These have 0-based indexing defined for them to match SymPy

The traditional infix mathmatical operations are defined, but no dot broadcasting.

The `convert(Matrix{Sym}, M)` call is useful to covert to a Julia matrix

"""
mutable struct SymMatrix <: SymbolicObject
    __pyobject__::PyCall.PyObject
end

## Permutations
## A permutation of {0, 1, 2, ..., n} -- 0-based
struct SymPermutation <: SymbolicObject
  __pyobject__::PyCall.PyObject
end
export SymPermutation
Base.convert(::Type{SymPermutation}, o::PyCall.PyObject) = SymPermutation(o)

## A permutation of {0, 1, 2, ..., n} -- 0-based
struct SymPermutationGroup <: SymbolicObject
  __pyobject__::PyCall.PyObject
end
export SymPermutationGroup
Base.convert(::Type{SymPermutationGroup}, o::PyCall.PyObject) = SymPermutationGroup(o)

# a Lambda function
struct Lambda <: SymbolicObject
    __pyobject__::PyCall.PyObject
end
Lambda(args, expression) = Lambda(sympy.Lambda(args, expression).__pyobject__)
(λ::Lambda)(args...; kwargs...) = λ.__pyobject__(args...; kwargs...)
export Lambda

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

## important override
## this allows most things to flow though PyCall
PyCall.PyObject(x::SymbolicObject) = x.__pyobject__
## Override this so that using symbols as keys in a dict works
hash(x::SymbolicObject) = hash(PyObject(x))
==(x::SymbolicObject, y::SymbolicObject) = PyObject(x) == PyObject(y)

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

## Show methods
"create basic printed output"
function jprint(x::SymbolicObject)
    out = PyCall.pycall(pybuiltin("str"), String, PyObject(x))
    out = replace(out, r"\*\*" => "^")
    out
end
jprint(x::AbstractArray) = map(jprint, x)

## text/plain
Base.show(io::IO, s::Sym) = print(io, jprint(s))
Base.show(io::IO, ::MIME"text/plain", s::SymbolicObject) =  print(io, sympy.pretty(s))

## latex enhancements: Sym, array, Dict
#Base.show(io::IO, ::MIME"text/latex", x::SymbolicObject) = print(io, sympy.latex(x, mode="equation*"))

as_markdown(x) = Markdown.parse("``$x``")
function Base.show(io::IO, ::MIME"text/latex", x::SymbolicObject)
    print(io, sympy.latex(x, mode="inline",fold_short_frac=false)) # plain? equation*?
#    #print(io, as_markdown(sympy.latex(x, mode="equation*")))
#    #print(io, as_markdown(sympy.latex(x, mode="plain",fold_short_frac=false))) # inline?
end
function  show(io::IO, ::MIME"text/latex", x::AbstractArray{Sym})
    function toeqnarray(x::Vector{Sym})
        a = join([sympy.latex(x[i]) for i in 1:length(x)], "\\\\")
        """\\left[ \\begin{array}{r}$a\\end{array} \\right]"""
#        "\\begin{bmatrix}$a\\end{bmatrix}"
    end
    function toeqnarray(x::AbstractArray{Sym,2})
        sz = size(x)
        a = join([join(map(sympy.latex, x[i,:]), "&") for i in 1:sz[1]], "\\\\")
        "\\left[ \\begin{array}{" * repeat("r",sz[2]) * "}" * a * "\\end{array}\\right]"
#        "\\begin{bmatrix}$a\\end{bmatrix}"
    end
    print(io, as_markdown(toeqnarray(x)))
end
function show(io::IO, ::MIME"text/latex", d::Dict{T,S}) where {T<:SymbolicObject, S<:Any}
    Latex(x::Sym) = sympy.latex(x)
    Latex(x) = sprint(io -> show(IOContext(io, :compact => true), x))

    out = "\\begin{equation*}\\begin{cases}"
    for (k,v) in d
        out = out * Latex(k) * " & \\text{=>} &" * Latex(v) * "\\\\"
    end
    out = out * "\\end{cases}\\end{equation*}"
    print(io, as_markdown(out))
end

latex(x::Sym) = sympy.latex(x)

# (this is used by Polynomials.jl)
function Base.show_unquoted(io::IO, pj::SymbolicObject, indent::Int, prec::Int)
    if Base.operator_precedence(:+) <= prec
        print(io, "(")
        show(io, pj)
        print(io, ")")
    else
        show(io, pj)
    end
end

## Following recent changes to PyCall where:
# For o::PyObject, make o["foo"], o[:foo], and o.foo equivalent to o.foo in Python,
# with the former returning an raw PyObject and the latter giving the PyAny
# conversion.
# We do something similar to SymPy
#
# We only implement for symbols here, not strings
function Base.getproperty(o::T, s::Symbol) where {T <: SymbolicObject}
    if (s in fieldnames(T))
        getfield(o, s)
    else
        getproperty(PyCall.PyObject(o), s)
    end
end

# XXX Needs version v1.2+
#function Base.hasproperty(o::T, s::Symbol) where {T <: SymbolicObject}
#    s ∈ fieldnames(T) && return true
#    hasproperty(PyCall.PyObject(o), s)
#end


text:
module GaussianRandomFields

# dependencies
using SpecialFunctions, RecipesBase, FastGaussQuadrature, Arpack, FFTW

using LinearAlgebra, Statistics

using Base.Cartesian

# export statements
export CovarianceStructure, IsotropicCovarianceStructure, AnisotropicCovarianceStructure, AbstractCovarianceFunction, CovarianceFunction, apply # from covariance_functions.jl

export SeparableCovarianceFunction # from separable.jl

export Matern # from matern.jl

export Exponential # from exponential.jl

export Linear # from linear.jl

export Spherical # from spherical.jl

export Whittle # from whittle.jl

export SquaredExponential, Gaussian # from squaredexponential.jl

export AnisotropicExponential # from anisotropic_exponential.jl

export GaussianRandomFieldGenerator, GaussianRandomField, sample, randdim # from gaussian_random_fields.jl

export Cholesky # from cholesky.jl

export Spectral # from spectral.jl

export KarhunenLoeve, rel_error # from karhunen_loeve.jl

export CirculantEmbedding # from circulant_embedding.jl

export QuadratureRule, GaussLegendre, EOLE, Simpson, Midpoint, Trapezoidal, AbstractEigenSolver, EigsSolver, EigenSolver # from quadrature.jl

export star, Lshape # from fem_data/

# include statements
include("covariance_functions/covariance_functions.jl")

include("covariance_functions/matern.jl")

include("covariance_functions/exponential.jl")

include("covariance_functions/linear.jl")

include("covariance_functions/spherical.jl")

include("covariance_functions/whittle.jl")

include("covariance_functions/squared_exponential.jl")

include("covariance_functions/anisotropic_exponential.jl")

include("gaussian_random_fields.jl")

include("fem.jl")

include("generators/cholesky.jl")

include("generators/spectral.jl")

include("generators/quadrature.jl")

include("generators/karhunen_loeve.jl")

include("generators/circulant_embedding.jl")

include("analytic.jl")

include("separable.jl")

include("../fem_data/star.jl")

include("../fem_data/Lshape.jl")

include("plot.jl")

end # module

ext: jl


text:
export creategaussian
export derive
export fouriertransform

immutable Gaussian{T}
    scaling::T
    width::T
    delay::T
end

(g::Gaussian)(s::Real) = 4*g.scaling/(g.width*√π) * exp(-(4*(s-g.delay)/g.width)^2)

function creategaussian(width,s0,scaling=one(typeof(width)))
    #f(s) = 4*scaling/(width*sqrt(π)) * exp(-(4*(s-s0)/width)^2)
    Gaussian(scaling, width, s0)
    #f(s) = scaling * exp(-(4*(s-s0)/width)^2)
end

function  fouriertransform(g::Gaussian)
    scaling = g.scaling
    width = g.width
    s0 = g.delay
    ft(w) = scaling * exp(-im*w*s0 - (width*w/8)^2) / sqrt(2π)
end

"""
```math
    F(w) = 1/\sqrt{2pi} \int f(t) e^{-i ω t} dt
```
"""
function fouriertransform(a::Array, dt, t0, dim=1)
    n = size(a,dim)
    dω = 2π / (n*dt)
    b = fftshift(fft(a, dim), dim) * dt / sqrt(2π)
    ω0 = -dω * div(n,2)
    b, dω, ω0
end

derive(g::Gaussian) =  s -> g(s) * (-8 * (s-g.delay)/g.width) * (4/g.width)


text:
module Model

using ParticleDA

using Random, Distributions, Base.Threads, GaussianRandomFields, HDF5
# using Default_params
using DelimitedFiles
using FieldMetadata
using FortranFiles
using Setfield
using Dates
using MPI

# using .Default_params
include("speedy.jl")
using .SPEEDY

include("grf.jl")
using .Periodic_Cov
"""
    ModelParameters()

Parameters for the model. Keyword arguments:

* `IDate::String` : Start date for the simulations in the format: YYYYmmddHH
* `dtDate::String` : Incremental date in the format: YYYYmmddHH
* `Hinc::Int` : Hourly increment
* `obs_network::String` : Location of observations (real or uniform)
* `nobs::Int` : Number of observation stations
* `lambda::AbstractFloat` : Length scale for Matérn covariance kernel in background noise
* `nu::AbstractFloat` : Smoothess parameter for Matérn covariance kernel in background noise
* `sigma::AbstractFloat` : Marginal standard deviation for Matérn covariance kernel in background noise
* `lambda_initial_state::AbstractFloat` : Length scale for Matérn covariance kernel in initial state of particles
* `nu_initial_state::AbstractFloat` : Smoothess parameter for Matérn covariance kernel in initial state of particles
* `sigma_initial_state::AbstractFloat` : Marginal standard deviation for Matérn covariance kernel in initial state of particles
* `padding::Int` : Min padding for circulant embedding gaussian random field generator
* `primes::Int`: Whether the size of the minimum circulant embedding of the covariance matrix can be written as a product of small primes (2, 3, 5 and 7). Default is `true`.
* `obs_noise_std::Float`: Standard deviation of noise added to observations of the true state
* `particle_dump_file::String`: file name for dump of particle state vectors
* `particle_dump_time::Int`: list of (one more more) time steps to dump particle states
* `SPEEDY::String` : Path to the SPEEDY directory
* `output_folder::String` : Output folder
* `station_filename::String` : Path to the station file which defines the observation locations
* `nature_dir::String`: Path to the directory with the ouputs of the nature run
* `observed_indices::Vector`: Vector containing the indices of the observed values in the state vector
* `nlon::Int`: Number of points in the longitude direction
* `nlat::Int`: Number of points in the latitude direction
* `lon_length::AbstractFloat`: Domain size in the lon direction
* `lat_length::AbstractFloat`: Domain size in the lat direction
* `dx::AbstractFloat` : Distance between grid points in the lon direction
* `dy::AbstractFloat` : Distance between grid points in the lat direction
* `nlev::Int`: Number of points in the vertical direction
* `n_state_var::Int`: Number of variables in the state vector
"""

Base.@kwdef struct ModelParameters{T<:AbstractFloat}
    # Initial date
    IDate::String=""
    # Incremental date
    dtDate::String=""
    # Hour interval
    Hinc::Int = 6
    # Ensemble initialisation
    ensDate::String=""
    # Choose observation network (choose "real" or "uniform")
    obs_network::String = "uniform"
    # Number of obs stations
    nobs::Int = 50

    lambda::Vector{T} = [1.0e0]#, 1.0e0, 1.0e0, 1.0e0, 1.0e0]
    nu::Vector{T} = [2.5]#, 2.5, 2.5, 2.5, 2.5]
    sigma::Vector{T} = [1000.0]#[0.1, 0.1, 0.1, 0.00001, 1000.0]

    lambda_initial_state::Vector{T} = [1.0e0]#, 1.0e0, 1.0e0, 1.0e0, 1.0e0]
    nu_initial_state::Vector{T} = [2.5]#, 2.5, 2.5, 2.5, 2.5]
    sigma_initial_state::Vector{T} = [1000.0]#[0.1, 0.1, 0.1, 0.0001, 1000.0]

    padding::Int = 100
    primes::Bool = true

    state_prefix::String = "data"
    title_avg::String = "avg"
    title_var::String = "var"
    title_syn::String = "syn"
    title_grid::String = "grid"
    title_stations::String = "stations"
    title_params::String = "params"
    particle_dump_file::String = "particle_dump.h5"
    particle_dump_time = [-1]
    obs_noise_std::T = 1000.0
    #Path to the the local speedy directory
    SPEEDY::String = ""
    station_filename::String = joinpath(SPEEDY, "obs", "networks", obs_network * ".txt")
    nature_dir::String = joinpath(SPEEDY, "DATA", "nature")
    # Output folders
    output_folder::String = joinpath(pwd())
    guess_folder::String = joinpath(output_folder, "DATA", "ensemble", "gues")
    anal_folder::String = joinpath(output_folder, "DATA", "ensemble", "anal")
    # Observed indices
    observed_indices::Vector{Int} = [1,5]
    # Grid dimensions
    nlon::Int = 96
    nlat::Int = 48
    nlev::Int = 8
    lon_length::T = 360.0
    lat_length::T = 180.0
    z_length::T = 1.0
    dx::T = lon_length / (nlon - 1)
    dy::T = lat_length / (nlat - 1)
    dz::T = z_length / (nlev - 1)

    nij0::Int = nlon*nlat
    n_state_var::Int = 6
    n_2d::Int = 2
    n_3d::Int = 4

end
const SPEEDY_DATE_FORMAT = "YYYYmmddHH"

 function step_datetime(idate::String,dtdate::String)
     new_idate = Dates.format(DateTime(idate, SPEEDY_DATE_FORMAT) + Dates.Hour(6), SPEEDY_DATE_FORMAT)
     new_dtdate = Dates.format(DateTime(dtdate, SPEEDY_DATE_FORMAT) + Dates.Hour(6), SPEEDY_DATE_FORMAT)
     return new_idate,new_dtdate
 end

 function step_ens(ensdate::String, Hinc::Int)
     ens_end = Dates.format(DateTime(ensdate, SPEEDY_DATE_FORMAT) + Dates.Month(1), SPEEDY_DATE_FORMAT)
     diff = length(DateTime(ensdate, SPEEDY_DATE_FORMAT):Hour(Hinc):DateTime(ens_end, SPEEDY_DATE_FORMAT))
     rand_int = rand(0:diff-1)*Hinc
     new_idate = Dates.format(DateTime(ensdate, SPEEDY_DATE_FORMAT) + Dates.Hour(rand_int), SPEEDY_DATE_FORMAT)
     return new_idate
 end

function ParticleDA.get_params(T::Type{ModelParameters}, user_input_dict::Dict)

    for key in ("lambda", "nu", "sigma", "lambda_initial_state", "nu_initial_state", "sigma_initial_state")
        if haskey(user_input_dict, key) && !isa(user_input_dict[key], Vector)
            user_input_dict[key] = fill(user_input_dict[key], 3)
        end
    end

    user_input = (; (Symbol(k) => v for (k,v) in user_input_dict)...)
    params = T(;user_input...)

end

function get_obs!(obs::AbstractVector{T},
                  state::AbstractArray{T},
                  ist::AbstractVector{Int},
                  jst::AbstractVector{Int},
                  params::ModelParameters) where T

    get_obs!(obs,state,ist,jst)

end

# Return observation data at stations from given model state
function get_obs!(obs::AbstractVector{T},
                  state::AbstractArray{T},
                  ist::AbstractVector{Int},
                  jst::AbstractVector{Int}) where T
    @assert length(obs) == length(ist) == length(jst)

    for i in eachindex(obs)
        ii = ist[i]
        jj = jst[i]
        obs[i] = state[ii,jj]
    end
end

struct RandomField{F<:GaussianRandomField,X<:AbstractArray,W<:AbstractArray,Z<:AbstractArray}
    grf::F
    xi::X
    w::W
    z::Z
end

struct ObsVectors{T<:AbstractArray,S<:AbstractArray}

    truth::T
    model::S

end

struct StationVectors{T<:AbstractArray}

    ist::T
    jst::T

end

@metadata name ("","","","","","") NTuple{6, String}
@metadata unit ("","","","","","") NTuple{6, String}
@metadata description ("","","","","","") NTuple{6, String}

@name @description @unit struct StateVectors{T<:AbstractArray, S<:AbstractArray}

    particles::T | ("u","v","T","q","ps","rain") | ("velocity x-component","velocity y-component","Temperature", "Specific Humidity", "Pressure", "Rain") | ("m/s","m/s","K","kg/kg","hPa","mm/hr")
    truth::S | ("ps") | ("Pressure") | ("hPa")

end

function get_axes(params::ModelParameters)

    return get_axes(params.nlon, params.nlat, params.nlev)

end
function get_axes(ix::Int, iy::Int, iz::Int)
    x = LinRange(0, 2*π, ix)
    y = LinRange(0, π, iy)
    z = LinRange(0, 1, iz)

    return x,y,z
end

function init_gaussian_random_field_generator(params::ModelParameters)

    x, y, z = get_axes(params)
    return init_gaussian_random_field_generator(params.lambda,params.nu, params.sigma, x, y, params.padding, params.primes)

end

# Initialize a gaussian random field generating function using the Matern covariance kernel
# and circulant embedding generation method
# TODO: Could generalise this
function init_gaussian_random_field_generator(lambda::Vector{T},
                                              nu::Vector{T},
                                              sigma::Vector{T},
                                              xpts::AbstractVector{T},
                                              ypts::AbstractVector{T},
                                              pad::Int,
                                              primes::Bool) where T

    # Let's limit ourselves to two-dimensional fields
    dim = 2
    function _generate(l, n, s)
        cov = Periodic_Cov.CovarianceFunction(dim, Periodic_Cov.CustomDistanceCovarianceStructure(Periodic_Cov.SphericalDistance(), Matern(l, n, σ = s)))
        grf = GaussianRandomField(cov, Spectral(), xpts, ypts)
        v = grf.data.eigenval
        xi = Array{eltype(grf.cov)}(undef, size(v)..., nthreads())
        w = Array{complex(float(eltype(grf.cov)))}(undef, size(v)..., nthreads())
        z = Array{eltype(grf.cov)}(undef, length.(grf.pts)..., nthreads())
        RandomField(grf, xi, w, z)
    end

    return [_generate(l, n, s) for (l, n, s) in zip(lambda, nu, sigma)]
end
# Get a random sample from random_field_generator using random number generator rng
function sample_gaussian_random_field!(field::AbstractArray{T,2},
                                       random_field_generator::RandomField,
                                       rng::Random.AbstractRNG) where T

    field .= GaussianRandomFields.sample(random_field_generator.grf)
end

# Add a gaussian random field to the height in the state vector of all particles
function add_random_field!(state::AbstractArray{T},
                           field_buffer::AbstractArray{T},
                           generators::Vector{<:RandomField},
                           rng::AbstractVector{<:Random.AbstractRNG},
                           n_3d::Int,
                           nlev::Int,
                           nprt::Int) where T

    Threads.@threads for ip in 1:nprt
        # for ivar in 1:n_3d
        #     # Adding model noise to the 3D fields (U,V,T,q)
        #     for lev in 1:nlev
        #         sample_gaussian_random_field!(@view(field_buffer[:, :, lev, ivar, threadid()]), generators[ivar], rng[threadid()])
        #     end
        # end
        # state[:, :, :, 1:n_3d, ip] .+= @view(field_buffer[:, :, :, 1:n_3d, threadid()])
        # Adding model noise to the surface pressure field
        sample_gaussian_random_field!(@view(field_buffer[:, :, 1, 5, threadid()]), generators[1], rng[threadid()])
        state[:, :, 1, 5, ip] .+= @view(field_buffer[:, :, 1, 5, threadid()])
    end

end

function add_noise!(vec::AbstractVector{T}, rng::Random.AbstractRNG, params::ModelParameters) where T

    add_noise!(vec, rng, 0.0, params.obs_noise_std)

end

# Add a (mean, std) normal distributed random number to each element of vec
function add_noise!(vec::AbstractVector{T}, rng::Random.AbstractRNG, mean::T, std::T) where T

    d = truncated(Normal(mean, std), 0.0, Inf)
    @. vec += rand((rng,), d)

end

function init_arrays(params::ModelParameters, nprt_per_rank)

    return init_arrays(params.nlon, params.nlat, params.nlev, params.n_state_var, params.nobs, nprt_per_rank)

end

function init_arrays(ix::Int, iy::Int, iz::Int, n_state_var::Int, nobs::Int, nprt_per_rank::Int)

    # TODO: ideally this will be an argument of the function, to choose a
    # different datatype.
    T = Float64

    state_avg = zeros(T, ix, iy, iz, n_state_var) # average of particle state vectors
    state_var = zeros(T, ix, iy, iz, n_state_var) # variance of particle state vectors

    state_particles = zeros(T, ix, iy, iz, n_state_var, nprt_per_rank)
    state_truth = zeros(T, ix, iy, 1) # model vector: true surface pressure (observation)
    obs_truth = Vector{T}(undef, nobs)          # observed
    obs_model = Matrix{T}(undef, nobs, nprt_per_rank) # forecasted

    # station location in digital grids
    ist = zeros(Int64,0)
    jst = zeros(Int64,0)

    # Buffer array to be used in the update
    field_buffer = Array{T}(undef, ix, iy, iz, 5, nthreads())
    return StateVectors(state_particles, state_truth), ObsVectors(obs_truth, obs_model), StationVectors(ist, jst), field_buffer
end

function set_initial_state!(states::StateVectors, model_matrices::SPEEDY.Matrices,
                            field_buffer::AbstractArray{T},
                            rng::AbstractVector{<:Random.AbstractRNG},
                            nprt_per_rank::Int,
                            params::ModelParameters) where T

    # Set true initial state
    # Read in the initial nature run - just surface pressure
    read_ps!(@view(states.truth[:,:,1]), joinpath(params.nature_dir, params.IDate * ".grd"), params.nlon, params.nlat, params.nlev)
    ### Read in arbitrary nature run files for the initial conditions
    Threads.@threads for ip in 1:nprt_per_rank
        dummy_date = step_ens(params.ensDate, params.Hinc)
        read_grd!(@view(states.particles[:,:,:,:,ip]), joinpath(params.nature_dir, dummy_date * ".grd"), params.nlon, params.nlat, params.nlev)
    end
end

# Set station locations.
function set_stations!(stations::StationVectors, params::ModelParameters) where T
    set_stations!(stations.ist,
                  stations.jst,
                  params.station_filename)
end

function set_stations!(ist::AbstractVector, jst::AbstractVector, filename::String)
    open(filename,"r") do f
        readline(f)
        readline(f)
        ind = 0
        for line in eachline(f)
            # if ind%10 == 0
            append!(ist, parse(Int64,split(line)[1]))
            append!(jst, parse(Int64,split(line)[2]))
            # end
            ind = ind + 1
        end
    end
end

struct ModelData{A,B,C,D,E,F,G,H,I}
    model_params::A
    states::B
    observations::C
    stations::D
    field_buffer::E
    background_grf::F
    model_matrices::G
    rng::H
    dates::I
end
ParticleDA.get_particles(d::ModelData) = d.states.particles
# TODO: we should probably get rid of `get_truth`: it is only used as return
# value of `particle_filter`, we may just return the whole `model_data`.
ParticleDA.get_truth(d::ModelData) = d.states.truth
ParticleDA.get_stations(d::ModelData) = (nst = d.model_params.nobs,
                                         ist = d.stations.ist,
                                         jst = d.stations.jst)
ParticleDA.get_obs_noise_std(d::ModelData) = d.model_params.obs_noise_std
# Needs to be corrected
ParticleDA.get_model_noise_params(d::ModelData) = Matern(d.model_params.lambda[1],
                                                         d.model_params.nu[1],
                                                         σ=d.model_params.sigma[1])
ParticleDA.get_observed_state_indices(d::ModelData) = d.model_params.observed_indices

function ParticleDA.set_particles!(d::ModelData, particles::AbstractArray{T}) where T

    d.states.particles .= particles

end
ParticleDA.get_grid_size(d::ModelData) = d.model_params.nlon,d.model_params.nlat,d.model_params.nlev
ParticleDA.get_grid_domain_size(d::ModelData) = d.model_params.lon_length,d.model_params.lat_length,d.model_params.z_length
ParticleDA.get_grid_cell_size(d::ModelData) = d.model_params.dx,d.model_params.dy,d.model_params.dz
ParticleDA.get_n_state_var(d::ModelData) = d.model_params.n_state_var

function create_folders(output_folder::String, anal_folder::String, gues_folder::String, nprt_per_rank::Int, my_rank::Integer)
    MPI.Init()
    comm = MPI.COMM_WORLD
    data = joinpath(output_folder, "DATA")
    ens = joinpath(data, "ensemble")
    tmp = joinpath(data, "tmp")
    if MPI.Comm_rank(comm) == 0
        rm(data; recursive=true)
        mkdir(data)
        mkdir(ens)
        mkdir(tmp)
        mkdir(anal_folder)
        mkdir(gues_folder)
    end
    MPI.Barrier(comm)
    rank_tmp = joinpath(tmp,string(my_rank))
    rank_anal = joinpath(anal_folder,string(my_rank))
    rank_gues = joinpath(gues_folder,string(my_rank))
    mkdir(rank_tmp)
    mkdir(rank_anal)
    mkdir(rank_gues)
    Threads.@threads for ip in 1:nprt_per_rank
        mkdir(joinpath(rank_anal, string(ip)))
        mkdir(joinpath(rank_gues, string(ip)))
        mkdir(joinpath(rank_tmp, string(ip)))
    end
end

function init(model_params_dict::Dict, nprt_per_rank::Int, my_rank::Integer, rng::Vector{<:Random.AbstractRNG})
    model_params = ParticleDA.get_params(ModelParameters, get(model_params_dict, "speedy", Dict()))
    states, observations, stations, field_buffer = init_arrays(model_params, nprt_per_rank)
    background_grf = init_gaussian_random_field_generator(model_params)
    # Set up model
    model_matrices = SPEEDY.setup(model_params.nlon,model_params.nlat, model_params.nlev)
    model_dates = [model_params.IDate,model_params.dtDate]
    set_stations!(stations, model_params)

    set_initial_state!(states, model_matrices, field_buffer, rng, nprt_per_rank, model_params)

    create_folders(model_params.output_folder, model_params.anal_folder, model_params.guess_folder, nprt_per_rank, my_rank)

    return ModelData(model_params, states, observations, stations, field_buffer, background_grf, model_matrices, rng, model_dates)
end
function read_grd!(truth::AbstractArray{T}, filename::String,nlon::Int, nlat::Int, nlev::Int) where T

    nij0 = nlon*nlat
    iolen = 4
    nv3d = 4
    nv2d = 2

    f = FortranFile(filename, access="direct", recl=nij0*iolen)

    irec = 1

    for n = 1:nv3d
        for k = 1:nlev
            truth[:,:,k,n] .= read(f, (Float32, nlon, nlat), rec=irec)
            irec += 1
        end
    end

    for n = 1:nv2d
        truth[:,:,1,(n+nv3d)] .= read(f, (Float32, nlon, nlat), rec = irec)
        irec += 1
    end

    close(f)

end

function read_ps!(truth::AbstractMatrix{T}, filename::String, nlon::Int, nlat::Int, nlev::Int) where T

    nij0 = nlon*nlat
    iolen = 4
    nv3d = 4
    nv2d = 2

    v3d = Array{Float32, 4}(undef, nlon, nlat, nlev, nv3d)
    v2d = Array{Float32, 3}(undef, nlon, nlat, nv2d)

    f = FortranFile(filename, access="direct", recl=nij0*iolen)

    irec = 1

    for n = 1:nv3d
        for k = 1:nlev
            v3d[:,:,k,n] = read(f, (Float32, nlon, nlat), rec=irec)
            irec += 1
        end
    end

    for n = 1:nv2d
        v2d[:,:,n] = read(f, (Float32, nlon, nlat), rec = irec)
        irec += 1
    end
    truth .= v2d[:,:,1]
    close(f)

end

function write_fortran(filename::String,nlon::Int, nlat::Int, nlev::Int,dataset::AbstractArray{T}) where T

    nij0 = nlon*nlat
    iolen = 4
    nv3d = 4
    nv2d = 2
    iolen = 4
    v3d = Array{Float32, 4}(undef, nlon, nlat, nlev, nv3d)
    v2d = Array{Float32, 3}(undef, nlon, nlat, nv2d)
    v3d .= dataset[:,:,:,1:4]
    v2d .= dataset[:,:,1,5:6]

    f = FortranFile(filename, "w", access="direct", recl=(nij0*iolen))
    irec = 1
    for n = 1:nv3d
        for k = 1:nlev
            write(f,v3d[:,:,k,n],rec=irec)
            irec=irec+1
        end
    end
    for n = 1:nv2d
        write(f,v2d[:,:,n],rec=irec)
        irec=irec+1
    end
    close(f)
end

function ParticleDA.update_truth!(d::ModelData, _)
    # Read in Nature run
    read_ps!(@view(d.states.truth[:,:,1]), joinpath(d.model_params.nature_dir, d.dates[1] * ".grd"), d.model_params.nlon, d.model_params.nlat, d.model_params.nlev)
    # Get observation from nature run
    get_obs!(d.observations.truth, d.states.truth[:,:,1], d.stations.ist, d.stations.jst, d.model_params)
    add_noise!(d.observations.truth, d.rng[1], d.model_params)
    return d.observations.truth
end

function speedy_update!(SPEEDY::String,
                        output::String,
                        YMDH::String,
                        TYMDH::String,
                        rank::String,
                        particle::String)
    # Path to the bash script which carries out the forecast
    forecast = joinpath(@__DIR__, "dafcst.sh")
    # Bash script call to speedy
    run(`$forecast $SPEEDY $output $YMDH $TYMDH $rank $particle`)
end

function ParticleDA.sample_observations_given_particles(d::ModelData, nprt_per_rank)
    indices = d.model_params.observed_indices
    for ip in 1:nprt_per_rank
        get_obs!(@view(d.observations.model[:,ip]),
                 @view(d.states.particles[:, :, indices..., ip]),
                 d.stations.ist,
                 d.stations.jst,
                 d.model_params)
        add_noise!(@view(d.observations.model[:,ip]), d.rng[threadid()], d.model_params)
    end
    return d.observations.model
end

function ParticleDA.update_particle_dynamics!(d::ModelData, nprt_per_rank)
    my_rank = MPI.Comm_rank(MPI.COMM_WORLD)
    Threads.@threads for ip in 1:nprt_per_rank
        #Write to file
        anal_file = joinpath(d.model_params.anal_folder, string(my_rank), string(ip), d.dates[1] * ".grd")
        write_fortran(anal_file,d.model_params.nlon, d.model_params.nlat, d.model_params.nlev,d.states.particles[:, :, :, :, ip])
        # Update the dynamics
        speedy_update!(d.model_params.SPEEDY,d.model_params.output_folder,d.dates[1],d.dates[2],string(my_rank),string(ip))
        # Read back in the data and update the states
        guess_file = joinpath(d.model_params.guess_folder, string(my_rank), string(ip), d.dates[2] * ".grd")
        read_grd!(@view(d.states.particles[:, :, :, :, ip]), guess_file, d.model_params.nlon, d.model_params.nlat, d.model_params.nlev)
        # Remove old files
        rm(anal_file)
        rm(guess_file)
    end
    ## Update the time strings
    d.dates[1],d.dates[2] = step_datetime(d.dates[1],d.dates[2])
    # @show d.dates[1], d.dates[2]
end

function ParticleDA.update_particle_noise!(d::ModelData, nprt_per_rank)
    # Add process noise
    add_random_field!(d.states.particles,
                      d.field_buffer,
                      d.background_grf,
                      d.rng,
                      d.model_params.n_3d,
                      d.model_params.nlev,
                      nprt_per_rank)
end

function ParticleDA.get_particle_observations!(d::ModelData, nprt_per_rank)
    indices = d.model_params.observed_indices
    # get observations
    for ip in 1:nprt_per_rank
        get_obs!(@view(d.observations.model[:,ip]),
                 @view(d.states.particles[:, :, indices...,  ip]),
                 d.stations.ist,
                 d.stations.jst,
                 d.model_params)
    end
    return d.observations.model
end

### Model IO

function write_params(output_filename, params)

    file = h5open(output_filename, "cw")

    if !haskey(file, params.title_params)

        group = create_group(file, params.title_params)

        fields = fieldnames(typeof(params));

        for field in fields

            attributes(group)[string(field)] = getfield(params, field)

        end

    else

        @warn "Write failed, group " * params.title_params * " already exists in " * file.filename * "!"

    end

    close(file)

end

function write_grid(output_filename, params)

    h5open(output_filename, "cw") do file

        if !haskey(file, params.title_grid)

            # Write grid axes
            group = create_group(file, params.title_grid)
            x,y,z = get_axes(params)
            #TODO: use d_write instead of create_dataset when they fix it in the HDF5 package
            ds_x,dtype_x = create_dataset(group, "x", collect(x))
            ds_y,dtype_x = create_dataset(group, "y", collect(y))
            ds_z,dtype_x = create_dataset(group, "z", collect(z))
            ds_x[1:params.nlon] = collect(x)
            ds_y[1:params.nlat] = collect(y)
            ds_z[1:params.nlev] = collect(z)
            attributes(ds_x)["Unit"] = "°"
            attributes(ds_y)["Unit"] = "°"
            attributes(ds_z)["Unit"] = "°"

        else

            @warn "Write failed, group " * params.title_grid * " already exists in " * file.filename * "!"

        end

    end

end

function write_stations(output_filename, ist::AbstractVector, jst::AbstractVector, params::ModelParameters) where T
    x,y,z = get_axes(params)
    h5open(output_filename, "cw") do file

        if !haskey(file, params.title_stations)
            group = create_group(file, params.title_stations)

            for (dataset_name, val) in zip(("x", "y"), (x[ist], y[jst]))
                ds, dtype = create_dataset(group, dataset_name, val)
                ds[:] = val
                attributes(ds)["Description"] = "Station "*dataset_name*" coordinate"
                attributes(ds)["Unit"] = "°"
            end
        else
            @warn "Write failed, group " * params.title_stations * " already exists in " * file.filename * "!"
        end
    end
end

function write_weights(file::HDF5.File, weights::AbstractVector, unit::String, it::Int, d::ModelData)

    group_name = "weights"
    dataset_name = "t" * lpad(string(it),4,'0')

    group, subgroup = ParticleDA.create_or_open_group(file, group_name)

    if !haskey(group, dataset_name)
        #TODO: use d_write instead of create_dataset when they fix it in the HDF5 package
        ds,dtype = create_dataset(group, dataset_name, weights)
        ds[:] = weights
        attributes(ds)["Description"] = "Particle Weights"
        attributes(ds)["Unit"] = unit
        attributes(ds)["Time step"] = it
        attributes(ds)["Date"] = d.dates[1]
    else
        @warn "Write failed, dataset " * group_name * "/" * dataset_name *  " already exists in " * file.filename * "!"
    end

end

function ParticleDA.write_snapshot(output_filename::AbstractString,
                                   d::ModelData,
                                   avg::AbstractArray{T,4},
                                   var::AbstractArray{T,4},
                                   weights::AbstractVector{T},
                                   it::Int) where T

    if it == 0
        # These are written only at the initial state it == 0
        write_grid(output_filename, d.model_params)
        write_params(output_filename, d.model_params)
        write_stations(output_filename, d.stations.ist, d.stations.jst, d.model_params)

    end

    if any(d.model_params.particle_dump_time .== it)
        write_particles(d.model_params.particle_dump_file, d.states, it, d)
    end

    return ParticleDA.write_snapshot(output_filename, d.states, avg, var, weights, it, d)
end

function ParticleDA.write_snapshot(output_filename::AbstractString,
                                   states::StateVectors,
                                   avg::AbstractArray{T,4},
                                   var::AbstractArray{T,4},
                                   weights::AbstractVector{T},
                                   it::Int,
                                   d::ModelData) where T

    println("Writing output at timestep = ", it)

    h5open(output_filename, "cw") do file
        for (i,(name,desc,unit)) in enumerate(zip(name(states, :particles), description(states, :particles), unit(states, :particles)))
            write_field(file, @view(states.truth[:,:,1]), it, unit, d.model_params.title_syn, name, desc, d)
            write_field(file, @view(avg[:,:,:,i]), it, unit, d.model_params.title_avg, name, desc, d)
            write_field(file, @view(var[:,:,:,i]), it, "("*unit*")^2", d.model_params.title_var, name, desc, d)

        end

        write_weights(file, weights, "", it, d)
    end

end

function write_particles(output_filename::AbstractString,
                         states::StateVectors,
                         it::Int,
                         d::ModelData) where T

    println("Writing particle states at timestep = ", it)
    nprt = size(states.particles,5)

    h5open(output_filename, "cw") do file

        for iprt = 1:nprt
            group_name = "particle" * string(iprt)

            for (i,(name,desc,unit)) in enumerate(zip(name(states, :particles), description(states, :particles), unit(states, :particles)))

                write_field(file, @view(states.particles[:,:,:,i,iprt]), it, unit, group_name, name, desc, d)

            end

        end

    end

end

function write_field(file::HDF5.File,
                     field::AbstractArray{T},
                     it::Int,
                     unit::String,
                     group::String,
                     dataset::String,
                     description::String,
                     d::ModelData) where T

    group_name = d.model_params.state_prefix * "_" * group
    subgroup_name = "t" * lpad(string(it),4,'0')
    dataset_name = dataset

    group, subgroup = ParticleDA.create_or_open_group(file, group_name, subgroup_name)

    if !haskey(subgroup, dataset_name)
        #TODO: use d_write instead of create_dataset when they fix it in the HDF5 package
        ds,dtype = create_dataset(subgroup, dataset_name, field)
        if ndims(field) < 3
            ds[:,:] = field
        else
            ds[:,:,:] = field
        end
        attributes(ds)["Description"] = description
        attributes(ds)["Unit"] = unit
        attributes(ds)["Time step"] = it
        attributes(ds)["Date"] = d.dates[1]
    else
        @warn "Write failed, dataset " * group_name * "/" * subgroup_name * "/" * dataset_name *  " already exists in " * file.filename * "!"
    end
end

end # module

max_stars_repo_name: Team-RADDISH/ParticleDA-UseCases

max_issues_repo_name: Team-RADDISH/ParticleDA-UseCases

max_forks_repo_name: Team-RADDISH/ParticleDA-UseCases


text:
function generateGraphNodeRanks!( grf::MetaDiGraph, isInit::Bool = true )
    fTree = generateFeasibleTree( grf, isInit )
    treeRanks = get_prop.( Ref(fTree), vertices(fTree), :rank )
    nsteps = 1
    maxsteps = 5 * nv(grf)

    while !( (leaveEdge = findLeaveEdge(fTree)) isa Nothing )
        # Get a new edge to replace the bad edge with.
        headNodes = breakFeasibleTree( SimpleGraph(fTree.graph), leaveEdge )
        validEdges = filter( edge -> ( edge.src ∈ headNodes ) &&
            ( edge.dst ∉ headNodes ), collect(edges(grf)) )
        relativeEdgeSlacks = map( validEdges) do edge
            ( treeRanks[edge.dst] - treeRanks[edge.src] ) -
                getEdgeMinSlack( grf, edge )
        end  # map(validEdges) do edge
        enterEdge = rand(validEdges[relativeEdgeSlacks .==
            minimum( relativeEdgeSlacks )])

        # Exchange the edges.
        clear_props!( fTree, leaveEdge )
        rem_edge!( fTree, leaveEdge )
        add_edge!( fTree, enterEdge )

        for prop in keys(props( grf, enterEdge ))
            add_prop!( fTree, enterEdge, prop,
                get_prop( grf, enterEdge, prop ) )
        end  # for prop in keys(props( grf, enterEdge ))

        generateCutValues!( fTree, grf )

        # Adjust the ranks of the nodes in the head of the tree.
        treeRanks[headNodes] .-= treeRanks[enterEdge.src] +
            getEdgeMinSlack( grf, enterEdge ) - treeRanks[enterEdge.dst]
        nsteps += 1
    end  # while !( (leaveEdge = findLeaveEdge(fTree)) isa Nothing )

    treeRanks .-= minimum(treeRanks)
    set_prop!.( Ref(grf), vertices(grf), :rank, treeRanks )
end  # generateGraphNodeRanks!( grf, isInit )

function generateFeasibleTree( grf::MetaDiGraph, isInit::Bool )
    # Generate initial node ranks if needed.
    if isInit || !all( has_prop.( Ref(grf), vertices(grf), :rank ) )
        generateInitialNodeRanks!(grf)
    end  # if isInit || ...

    graphEdges = collect(edges(grf))
    treeRanks = zeros( Float64, nv(grf) )
    fTree = MetaDiGraph(nv(grf) )

    # Add a random start node to the feasible tree.
    inTree = [rand(vertices(grf))]

    # In each step, add an edge from/to a non-tree node which has minimal rank
    #   difference (slack) in the original graph.
    while length(inTree) < nv(grf)
        # Find all edges incident on the current feasible tree.
        validEdges = filter( edge -> xor( edge.src ∈ inTree,
            edge.dst ∈ inTree ), graphEdges )
        edgeSlacks = map( edge -> get_prop( grf, edge.dst, :rank ) -
            get_prop( grf, edge.src, :rank ), validEdges )

        # Get all tight edges among incident edges.
        isTight = edgeSlacks .== getEdgeMinSlack.( Ref(grf), validEdges )
        !any(isTight) && (isTight = edgeSlacks .== minimum(edgeSlacks))  # Insurance!
        tightEdges = validEdges[isTight]
        tightSlacks = edgeSlacks[isTight]

        # If there are no candidates (bad initial ranks), re-initialise the
        #   node ranks and restart.
        if isempty(tightEdges)
            generateInitialNodeRanks!(grf)
            inTree = [rand(vertices(grf))]
            continue
        end  # if isempty( tightEdges )

        # Select tight edge with minimal slack.
        minSlack = minimum(tightSlacks)
        newEdge = rand(tightEdges[tightSlacks .== minSlack])

        # Find the new node and its rank in the tree.
        isNewDst = newEdge.src ∈ inTree
        treeNode = getfield( newEdge, isNewDst ? :src : :dst )
        newNode = getfield( newEdge, isNewDst ? :dst : :src )
        treeRanks[newNode] = treeRanks[treeNode] +
            getEdgeMinSlack( grf, newEdge ) * ( isNewDst ? 1 : -1 )

        # Add the new edge.
        add_edge!( fTree, newEdge )
        setEdgeWeight!( fTree, newEdge, getEdgeWeight( grf, newEdge ) )
        setEdgeMinSlack!( fTree, newEdge, getEdgeMinSlack( grf, newEdge ) )
        push!( inTree, newNode )
    end  # while length(inTree) < nv(grf)

    # Set all node ranks in the feasible tree.
    treeRanks .-= minimum(treeRanks)
    set_prop!.( Ref(fTree), vertices(fTree), :rank, treeRanks )

    # Generate cut values for each edge in the feasible tree.
    generateCutValues!( fTree, grf )
    fTree
end  # generateFeasibleTree( grf, isInit )

function generateInitialNodeRanks!( grf::MetaDiGraph )
    ranks = - ones( Float64, nv( grf ) )
    unprocessedClearNodes = Queue{Int}()

    # In each step, get the list of nodes that no unprocessed in neighbours,
    #   and assign them the lowest possible rank which respects the minimum
    #   slacks for each incoming edge.
    while any( ranks .< 0 )
        getClearNodes( grf, ranks, unprocessedClearNodes )
        node = dequeue!(unprocessedClearNodes)

        minRanks = map(inneighbors( grf, node )) do srcNode
            ranks[srcNode] + getEdgeMinSlack( grf, srcNode, node )
        end  # map(inneighbors( grf, node )) do srcNode

        ranks[node] = isempty(minRanks) ? 0 : maximum(minRanks)
    end  # while any( ranks .< 0 )

    set_prop!.( Ref(grf), vertices(grf), :rank, ranks )
end  # generateInitialNodeRanks!( grf )

function getClearNodes( grf::MetaDiGraph, ranks::Vector{Float64},
    unprocessed::Queue{Int} )
    for node in filter( node -> ranks[node] < 0, vertices(grf) )
        if ( node ∉ unprocessed ) &&
            all( ranks[inneighbors( grf, node )] .>= 0 )
            enqueue!( unprocessed, node )
        end  # if ( node ∉ unprocessed ) && ...
    end  # for node in filter( node -> ranks[node] < 0, vertices(grf) )
end  # getClearNodes( grf, ranks, unprocessed )

function generateCutValues!( fTree::MetaDiGraph, grf::MetaDiGraph )
    uTree = SimpleGraph(fTree.graph)
    grfEdges = collect(edges(grf))

    for edge in edges(fTree)
        headNodes = breakFeasibleTree( uTree, edge )
        cutEdges = filter( tmpEdge -> xor( tmpEdge.src ∈ headNodes,
            tmpEdge.dst ∈ headNodes ), grfEdges )
        isPosEdge = map( tmpEdge -> tmpEdge.dst ∈ headNodes, cutEdges )
        cutVal = sum(getEdgeWeight.( Ref(grf), cutEdges[isPosEdge] ))

        if !all(isPosEdge)
            cutVal -= sum(getEdgeWeight.( Ref(grf), cutEdges[.!isPosEdge] ))
        end  # if !all(isPosEdge)

        set_prop!( fTree, edge, :cut, cutVal )
    end  # for edge in edges(fTree)
end  # generateCutValues!( fTree, grf )

function findLeaveEdge( fTree::MetaDiGraph )
    negCutEdges = filter( edge -> get_prop( fTree, edge, :cut ) < 0,
        collect(edges(fTree)) )
    isempty(negCutEdges) ? nothing : rand(negCutEdges)
end  # findLeaveEdge( fTree )

function getEdgeMinSlack( grf::MetaDiGraph, edge::Edge )
    !has_edge( grf, edge ) && return 0
    has_prop( grf, edge, :minSlack ) ?
        get_prop( grf, edge, :minSlack ) : 1
end  # getEdgeMinSlack( grf, edge )

getEdgeMinSlack( grf::MetaDiGraph, src::Int, dst::Int ) =
    getEdgeMinSlack( grf, Edge( src, dst ) )
getEdgeMinSlack( grf::MetaDiGraph, edge::Tuple{Int, Int} ) =
    getEdgeMinSlack( grf, Edge(edge) )

function setEdgeMinSlack!( grf::MetaDiGraph, edge::Edge, minSlack::Real )
    ( !has_edge( grf, edge ) || ( minSlack < 0.0 ) ) && return
    set_prop!( grf, edge, :minSlack, Float64(minSlack) )
end  # setEdgeMinSlack!( grf, edge, minSlack )

setEdgeMinSlack!( grf::MetaDiGraph, src::Int, dst::Int, minSlack::Real ) =
    setEdgeMinSlack!( grf, Edge( src, dst ),  minSlack )
setEdgeMinSlack!( grf::MetaDiGraph, edge::Tuple{Int, Int}, minSlack::Real ) =
    setEdgeMinSlack!( grf, Edge(edge), minSlack )

function getEdgeWeight( grf::MetaDiGraph, edge::Edge )
    !has_edge( grf, edge ) && return 0.0
    !has_prop( grf, edge, weightfield(grf) ) && return grf.defaultweight
    get_prop( grf, edge, weightfield(grf) )
end  # getEdgeWeight( grf, edge )

getEdgeWeight( grf::MetaDiGraph, src::Int, dst::Int ) =
    getEdgeWeight( grf, Edge( src, dst ) )
getEdgeWeight( grf::MetaDiGraph, edge::Tuple{Int, Int} ) =
    getEdgeWeight( grf, Edge(edge) )

function setEdgeWeight!( grf::MetaDiGraph, edge::Edge, weight::Real )
    ( !has_edge( grf, edge ) || ( weight < 0.0 ) ) && return
    set_prop!( grf, edge, weightfield(grf), Float64(weight) )
end  # setEdgeWeight!( grf, edge, weight )

setEdgeWeight!( grf::MetaDiGraph, src::Int, dst::Int, weight::Real ) =
    setEdgeWeight!( grf, Edge( src, dst ),  weight )
setEdgeWeight!( grf::MetaDiGraph, edge::Tuple{Int, Int}, weight::Real ) =
    setEdgeWeight!( grf, Edge(edge), weight )

function breakFeasibleTree( uTree::SimpleGraph, edge::Edge )
    ii = 1
    # We have to remember that the edge is defined as tail-to-head (src-to-dst).
    headNodes = [edge.dst]

    while ii <= length(headNodes)
        newNodes = neighbors( uTree, headNodes[ii] )

        if ii == 1
            append!( headNodes, newNodes[newNodes .!= edge.src] )
        else
            append!( headNodes, filter( node -> node ∉ headNodes, newNodes ) )
        end  # if ii == 1

        ii += 1
    end  # while ii <= length(headNodes)

    headNodes
end  # breakFeasibleTree( uTree, edge )


text:
using Fortran
using Test

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


text:
@muladd function calc_twopoint_random!(_dW, integrator, dW)
    @.. _dW = ifelse(sign(dW) > 0.0, integrator.sqdt, -integrator.sqdt)
end

@muladd function calc_twopoint_random(integrator, dW)
  return sign(dW) > 0.0 ? integrator.sqdt : -integrator.sqdt
end

@muladd function calc_threepoint_random!(_dW, integrator, quantile, dW_scaled)
  for (index, value) in enumerate(dW_scaled)
    if value < quantile
      _dW[index] = -sqrt(3*integrator.dt)
    elseif value > -quantile
      _dW[index] = sqrt(3*integrator.dt)
    else
      _dW[index] = zero(integrator.dt)
    end
  end
end

@muladd function calc_threepoint_random(integrator, quantile, dW_scaled)
  if dW_scaled < quantile
    _dW = -sqrt(3*integrator.dt)
  elseif dW_scaled > -quantile
    _dW = sqrt(3*integrator.dt)
  else
    _dW = zero(integrator.dt)
  end
  _dW
end


text:
@testset "scaledkernel" begin
    rng = MersenneTwister(123456)
    x = randn(rng)
    y = randn(rng)
    s = rand(rng) + 1e-3

    k = SqExponentialKernel()
    ks = ScaledKernel(k, s)
    @test ks(x, y) == s * k(x, y)
    @test ks(x, y) == (s * k)(x, y)

    # Standardised tests.
    TestUtils.test_interface(k, Float64)
    test_ADs(x -> exp(x[1]) * SqExponentialKernel(), rand(1))

    test_params(s * k, (k, [s]))
end


text:
using StatsBase

"""
   `pegasos(x, y; <keyword arguments>)`

Solves the linear SVM problem on the given training set of examples `x`
and labels `y` with the Pegasos (primal estimated sub-gradient solver)
method by Shai Shalev-Shwartz, Yoram Singer, Nathan Srebro and Andrew Cotter.

### Keyword arguments
* lambda=0.05: a step size parameter
* iterations=1000000: the number of iterations performed until `w` is returned
* batchsize=1: The number of data points used in each iteration
* projection=false: In each iteration the `w` vector is projected onto the ball of radius 1/sqrt(lambda)

### Returns
* weight vector `w`
"""
function pegasos(x::Array{Float64,2}, y::Array{Float64,1}; lambda=0.05, iterations=1000000, batchsize=1, projection=false)
    tic()

    s,d = size(x)
    w = zeros(Float64,d) # start with w as the zero vector

    for t in 1:iterations
        # compute step size
        stepsize = 1.0/(lambda*t)

        # draw batchsize indices at random
        indices = sample(1:s,batchsize,replace=false)

        w_new = 0
        for i in indices
            xi = vec(x[i,:])
            if y[i]*vecdot(w,xi) < 1
                w_new += y[i]*xi
            end
        end
        w_new *= stepsize/float(batchsize)
        w_new += (1-1.0/t)*w
        w = w_new

        # perform the optional projection step
        # to the ball of radius 1/sqrt(lambda):
        if projection
            w = min(1,1/(sqrt(lambda)*norm(w)))*w
        end

    end

    @printf "finished after %d iterations and %.3f seconds\n" iterations toq()
    return w
end


text:
module CMBuffers

import CUDA

using MPI
using KernelAbstractions
using StaticArrays

export CMBuffer
export SingleCMBuffer, DoubleCMBuffer
export get_stage, get_transfer, prepare_transfer!, prepare_stage!

@enum CMBufferKind begin
    SingleCMBuffer
    DoubleCMBuffer
end

device(::Union{Array, SArray, MArray}) = CPU()
device(::CUDA.CuArray) = CUDADevice()

###
# Note: We use pinned and device-mapped hostbuffers in double staging
# Potential improvements
# - Investigate if we need to device-map
# - Implement single buffered with pinned + device-mapped hostbuffers

"""
    CMBuffer{T}(::Type{Arr}, kind, dims...; pinned = true)

CUDA/MPI buffer abstracts storage for MPI communication. The buffer is
used for staging data and for MPI transfers. When running on:

  - CPU -- a single buffer is used for staging and MPI transfers can be
    initiated directly to/from it.
  - CUDA -- either:
    - MPI is CUDA-aware: a single buffer on the device for staging and MPI
      transfers can be initiated directly to/from it, or
    - MPI is not CUDA-aware: a double buffering scheme with the staging
      buffer on the device and a transfer buffer on the host

# Arguments
- `T`: element type
- `Arr::Type`: what kind of array to allocate for `stage`
- `kind::CMBufferKind`: either `SingleCMBuffer` or `DoubleCMBuffer`
- `dims...`: dimensions of the array
"""
struct CMBuffer{T, Arr, Buff}
    stage::Arr     # Same type as Q.data, used for staging
    transfer::Buff # Union{Nothing,Buff}

    function CMBuffer{T}(::Type{Arr}, kind, dims...) where {T, Arr}
        if kind == SingleCMBuffer
            transfer = nothing
        elseif kind == DoubleCMBuffer
            transfer = zeros(T, dims...)
        else
            error("CMBufferkind $kind is not implemented yet")
        end

        stage = similar(Arr, T, dims...)
        buffer = new{T, typeof(stage), typeof(transfer)}(stage, transfer)

        return buffer
    end
end

function get_stage(buf::CMBuffer)
    return buf.stage
end

function get_transfer(buf::CMBuffer)
    if buf.transfer === nothing
        return buf.stage
    else
        return buf.transfer
    end
end

function prepare_transfer!(
    buf::CMBuffer;
    dependencies = nothing,
    progress = yield,
)
    if buf.transfer === nothing || length(buf.transfer) == 0
        return MultiEvent(dependencies)
    end

    event = async_copy!(
        device(buf.stage),
        buf.transfer,
        buf.stage;
        dependencies = dependencies,
        progress = progress,
    )

    return event
end

function prepare_stage!(buf::CMBuffer; dependencies = nothing, progress = yield)
    if buf.transfer === nothing || length(buf.stage) == 0
        return MultiEvent(dependencies)
    end

    event = async_copy!(
        device(buf.stage),
        buf.stage,
        buf.transfer;
        dependencies = dependencies,
        progress = progress,
    )

    return event
end

end # module


text:
#===============================================================================
# Some tests for my new linked list type.
#
# Author: Carl Cortright
# Date: 1/1/2016
#
===============================================================================#
# Clear the workspace to avoid conflicts
workspace()
include("linkedlist.jl")

using linkedlist

ll = DoubleLinkedList{String}()

println("Time to add a bunch of stuff!!!")
append(ll, "Hello world!")
append(ll, "Foo")
append(ll, "Bar")
append(ll, "Making")
append(ll, "a")
append(ll, "dependent")
append(ll, "clause.")

for i in 1:listsize(ll)
  println(access(ll, i))
end

lenOfll = listsize(ll)
println("The length of the list is $lenOfll")

println(remove(ll, 6))
println(remove(ll, 4))
insert(ll, "I", 4)
insert(ll, "am", 5)
insert(ll, "making", 6)
insert(ll, "independent", 8)

for i in 4:listsize(ll)
  print(access(ll, i))
  print(" ")
end

max_issues_repo_name: carlcortright/Playground


text:
export SMatSpace, sparse_matrix, nnz, sparsity, density

################################################################################
#
#  Parent constructor
#
################################################################################

function SMatSpace(R::Ring, r::Int, c::Int; cached = true)
  T = elem_type(R)
  return SMatSpace{T}(R, r, c, cached)
end

################################################################################
#
#  Basic functions
#
################################################################################

base_ring(A::SMatSpace{T}) where {T} = A.base_ring::parent_type(T)

parent(A::SMat) = SMatSpace(base_ring(A), A.r, A.c)

base_ring(A::SMat{T}) where {T} = A.base_ring::parent_type(T)

@doc Markdown.doc"""
    nrows(A::SMat) -> Int

Return the number of rows of $A$.
"""
function nrows(A::SMat)
  @assert A.r == length(A.rows)
  return A.r
end

@doc Markdown.doc"""
    ncols(A::SMat) -> Int

Return the number of columns of $A$.
"""
function ncols(A::SMat)
  return A.c
end

@doc Markdown.doc"""
    nnz(A::SMat) -> Int

Return the number of non-zero entries of $A$.
"""
function nnz(A::SMat)
  return A.nnz
end

size(A::SMat) = (nrows(A), ncols(A))

################################################################################
#
#  Hashing
#
################################################################################

function hash(A::SMat, h::UInt)
  return hash(A.r, hash(A.c, hash(A.rows, h)))
end

################################################################################
#
#  Comparison
#
################################################################################

@doc Markdown.doc"""
    ==(A::SMat, B::SMat) -> Bool

Return whether $A$ is equal to $B$.
"""
function ==(x::SMat{T}, y::SMat{T}) where T
  return base_ring(x) == base_ring(y) && x.rows == y.rows && ncols(x) == ncols(y)
end

################################################################################
#
#  Sparsity and density
#
################################################################################

@doc Markdown.doc"""
    sparsity(A::SMat) -> Float64

Return the sparsity of `A`, that is, the number of zero-valued elements divided
by the number of all elements.
"""
function sparsity(A::SMat)
  return 1.0 - nnz(A)/(nrows(A) * ncols(A))
end

@doc Markdown.doc"""
    density(A::SMat) -> Float64

Return the density of `A`, that is, the number of nonzero-valued elements
divided by the number of all elements.
"""
density(A::SMat) = 1.0 - sparsity(A)

################################################################################
#
#  String I/O
#
################################################################################

function show(io::IO, A::SMat{T}) where T
  print(io, "Sparse ", A.r, " x ", A.c, " matrix with ")
  print(io, A.nnz, " non-zero entries\n")
end

################################################################################
#
#  Empty sparse matrix (given base ring)
#
################################################################################

@doc Markdown.doc"""
    sparse_matrix(R::Ring) -> SMat

Return an empty sparse matrix with base ring $R$.
"""
function sparse_matrix(R::T) where T <: Ring
  r = SMat{elem_type(R)}()
  r.base_ring = R
  return r
end

################################################################################
#
#  Copy
#
################################################################################

function copy(A::SMat{T}) where T
  return sub(A, 1:nrows(A), 1:ncols(A))
end

################################################################################
#
#  Entry/row access
#
################################################################################

@doc Markdown.doc"""
    getindex(A::SMat, i::Int, j::Int)

Given a sparse matrix $A = (a_{ij})_{i, j}$, return the entry $a_{ij}$.
"""
function getindex(A::SMat{T}, i::Int, j::Int) where T
  if i in 1:A.r
    ra = A.rows[i]
    p = findfirst(isequal(j), ra.pos)
    if !(p === nothing)
      return ra.values[p]
    end
  end
  return zero(base_ring(A))
end

@doc Markdown.doc"""
    getindex(A::SMat, i::Int) -> SRow

Given a sparse matrix $A$ and an index $i$, return the $i$-th row of $A$.
"""
function getindex(A::SMat{T}, i::Int) where T
  (i < 1 || i > nrows(A)) && error("Index must be between 1 and $(nrows(A))")
  return A.rows[i]
end

@doc Markdown.doc"""
    setindex!(A::SMat, b::SRow, i::Int)

Given a sparse matrix $A$, a sparse row $b$ and an index $i$, set the $i$-th
row of $A$ equal to $b$.
"""
function setindex!(A::SMat{T}, b::SRow{T}, i::Int) where T
  (i < 1 || i > nrows(A)) && error("Index must be between 1 and $(nrows(A))")
  A.nnz = A.nnz - length(A.rows[i].pos) + length(b.pos)
  A.rows[i] = b
  return A
end

################################################################################
#
#  Random row from sparse matrix
#
################################################################################

@doc Markdown.doc"""
    rand_row(A::SMat) -> SRow

Return a random row of the sparse matrix $A$.
"""
function rand_row(A::SMat{T}) where T
  return rand(A.rows)
end

################################################################################
#
#  Sparse matrix constructors
#
################################################################################

@doc Markdown.doc"""
    sparse_matrix(A::MatElem; keepzrows::Bool = true)

Constructs the sparse matrix corresponding to the dense matrix $A$. If
`keepzrows` is false, then the constructor will drop any zero row of $A$.
"""
function sparse_matrix(A::MatElem; keepzrows::Bool = true)
  R = base_ring(A)
  m = sparse_matrix(R)
  m.c = ncols(A)
  m.r = 0

  for i=1:nrows(A)
    if iszero_row(A, i)
      if !keepzrows
        continue
      else
        r = SRow{elem_type(R)}()
      end
    else
      r = SRow{elem_type(R)}()
      for j = 1:ncols(A)
        t = A[i, j]
        if t != 0
          m.nnz += 1
          push!(r.values, R(t))
          push!(r.pos, j)
        end
      end
    end
    push!(m.rows, r)
    m.r += 1
  end
  return m
end

@doc Markdown.doc"""
    sparse_matrix(A::Array{T, 2}) -> SMat{T}

Constructs the sparse matrix corresponding to $A$.
"""
function sparse_matrix(A::Array{T, 2}) where {T <: RingElement}
  length(A) == 0 && error("Cannot create sparse matrix from empty array")
  m = sparse_matrix(parent(A[1, 1]))
  m.c = Base.size(A, 2)
  m.r = 0
  for i in 1:size(A, 1)
    if iszero_row(A, i)
      push!(m, SRow{T}())
      continue
    end
    r = SRow{T}()
    for j in 1:size(A, 2)
      if !iszero(A[i, j])
        m.nnz += 1
        push!(r.values, A[i, j])
        push!(r.pos, j)
      end
    end
    push!(m.rows, r)
    m.r += 1
  end
  return m
end

@doc Markdown.doc"""
    sparse_matrix(R::Ring, A::Array{T, 2}) -> SMat

Constructs the sparse matrix over $R$ corresponding to $A$.
"""
function sparse_matrix(R::Ring, A::Array{T, 2}) where {T}
  B = convert(Array{elem_type(R), 2}, map(R, A))
  return sparse_matrix(B)
end

function (M::SMatSpace)(A::T; R::S = base_ring(A),
                                 keepzrows::Bool = true) where {T <: MatElem, S <: Ring}
  return sparse_matrix(A, R, keepzrows)
end

function (M::SMatSpace)(A::Array{T, 2}) where T <: MatElem
  return sparse_matrix(A)
end

################################################################################
#
#  Modular reduction
#
################################################################################

@doc Markdown.doc"""
    mod_sym!(A::SMat{fmpz}, n::fmpz)

Inplace reduction of all entries of $A$ modulo $n$ to the symmetric residue
system.
"""
function mod_sym!(A::SMat{fmpz}, b::fmpz)
  for r in A
    mod_sym!(r, b)
  end
  return A
end

################################################################################
#
#  Ring change
#
################################################################################

@doc Markdown.doc"""
    map_entries(f, A::SMat) -> SMat

Given a sparse matrix $A$ and a callable object $f$, this function will
construct a new sparse matrix by applying $f$ to all elements of $A$.
"""
function map_entries(f, A::SMat{T}) where {T}
  x = zero(base_ring(A))
  y = f(x)
  S = parent(y)
  z = sparse_matrix(S)
  z.c = ncols(A)
  z.nnz = 0
  for r in A
    if iszero(r)
      push!(z, sparse_row(S))
    else
      rf = map_entries(f, r)
      push!(z, rf)
    end
  end
  return z
end

@doc Markdown.doc"""
    change_base_ring(R::Ring, A::SMat)

Create a new sparse matrix by coercing all elements into the ring $R$.
"""
function change_base_ring(R::Ring, A::SMat)
  z = sparse_matrix(R)
  z.c = A.c
  for r in A
    rz = change_base_ring(R, r)
    push!(z, rz)
  end
  return z
end

################################################################################
#
#  Transpose
#
################################################################################

@doc Markdown.doc"""
    transpose(A::SMat) -> SMat

Returns the transpose of $A$.
"""
function transpose(A::SMat{T}) where {T}
  R = base_ring(A)
  B = sparse_matrix(R)
  n = nrows(A)
  m = ncols(A)
  B.rows = Vector{SRow{T}}(undef, m)
  for i=1:m
    B.rows[i] = SRow(R)
  end
  for i = 1:length(A.rows)
    for j = 1:length(A.rows[i].pos)
      r = A.rows[i].pos[j]
      push!(B.rows[r].pos, i)
      push!(B.rows[r].values, A.rows[i].values[j])
    end
  end
  B.c = n
  B.r = m
  B.nnz = A.nnz
  return B
end

################################################################################
#
#  Make sparse matrices iteratable
#
################################################################################

function endof(A::SMat)
  return length(A.rows)
end

function Base.iterate(A::SMat, st::Int = 1)
  if st > nrows(A)
    return nothing
  end

  return A.rows[st], st + 1
end

function length(A::SMat)
  return nrows(A)
end

Base.eltype(A::SMat{T}) where {T} = SRow{T}

################################################################################
#
#  Multiplication
#
################################################################################

# SMat{T} * Vector{T} as (dense Array{T, 1})
# inplace
function mul!(c::Array{T, 1}, A::SMat{T}, b::AbstractArray{T, 1}) where T
  R = base_ring(A)
  z = zero(R)
  for (i, r) in enumerate(A)
    c[i] = dot(r, b, z)
  end
  return c
end

# (dense Array{T, 1}) * SMat{T} as (dense Array{T, 1})
@doc Markdown.doc"""
    mul(A::SMat{T}, b::AbstractVector{T}) -> Vector{T}

Return the product $A \cdot b$ as a dense vector.
"""
function mul(A::SMat{T}, b::AbstractVector{T}) where T
  @assert length(b) == ncols(A)
  c = Vector{T}(undef, nrows(A))
  mul!(c, A, b)
  return c
end

function mul!(c::Array{T, 2}, A::SMat{T}, b::AbstractArray{T, 2}) where T
  sz = size(b)
  @assert sz[1] == ncols(A)
  tz = size(c)
  @assert tz[1] == nrows(A)
  @assert tz[2] == sz[2]
  z = zero(base_ring(A))
  for m in 1:size(b, 2)
    for i in 1:nrows(A)
      c[i, m] = dot(A.rows[i], view(b, :, m), z)
    end
  end
  return c
end

# - SMat{T} * Array{T, 2} as Array{T, 2}
@doc Markdown.doc"""
    mul(A::SMat{T}, b::AbstractArray{T, 2}) -> Array{T, 2}

Return the product $A \cdot b$ as a dense array.
"""
function mul(A::SMat{T}, b::AbstractArray{T, 2}) where T
  sz = size(b)
  @assert sz[1] == ncols(A)
  c = Array{T}(undef, sz[1], sz[2])
  return mul!(c, A, b)
end

# - SMat{T} * MatElem as MatElem
# - Inplace
function mul!(c::MatElem{T}, A::SMat{T}, b::MatElem{T}) where T
  @assert nrows(b) == ncols(A)
  @assert nrows(c) == nrows(A)
  @assert ncols(c) == ncols(b)
  for m = 1:ncols(b)
    for i = 1:nrows(A)
      c[i, m] = dot(A.rows[i], b, m)
    end
  end
  return c
end

# - SMat{T} * MatElem{T} as MatElem{T}

@doc Markdown.doc"""
    mul(A::SMat{T}, b::MatElem{T}) -> MatElem

Return the product $A \cdot b$ as a dense matrix.
"""
function mul(A::SMat{T}, b::MatElem{T}) where T
  @assert nrows(b) == ncols(A)
  c = similar(b, nrows(A), ncols(b))
  return mul!(c, A, b)
end

# - SRow{T} * SMat{T} as SRow{T}

@doc Markdown.doc"""
    mul(A::SRow, B::SMat) -> SRow

Return the product $A\cdot B$ as a sparse row.
"""
function mul(A::SRow{T}, B::SMat{T}) where T
  C = sparse_row(base_ring(B))
  for (p, v) in A
    C = add_scaled_row(B[p], C, v)
  end
  return C
end

################################################################################
#
#  Multiplication with reduction
#
################################################################################

# - (dense Array{S, 1}) * SMat{T} as (dense Array{S, 1}) modulo mod::S
# - Inplace
# - Reduction as the last step, no intermediate reductions.
function mul_mod!(c::Array{S, 1}, A::SMat{T}, b::Array{S, 1}, mod::S) where {S, T}
  @assert length(b) == ncols(A)
  @assert length(c) == nrows(A)

  for i = 1:length(A.rows)
    s = S(0)
    I = A.rows[i]
    for j=1:length(I.pos)
      s += S(I.values[j]) * b[I.pos[j]]
    end
    c[i] = s % mod
  end
  return c
end

# - (dense Array{S, 1}) * SMat{T} as (dense Array{S, 1}) modulo mod::S
# - Inplace
# - Intermediate reductions.
function mul_mod_big!(c::Array{S, 1}, A::SMat{T}, b::Array{S, 1}, mod::S) where {S, T}
  @assert length(b) == ncols(A)
  @assert length(c) == nrows(A)
  for i = 1:length(A.rows)
    s = 0
    I = A.rows[i]
    for j=1:length(I.pos)
      s += S(I.values[j] *b[I.pos[j]] % mod)
    end
    c[i] = s % mod
  end
  return c
end

################################################################################
#
#  Addition
#
################################################################################

@doc Markdown.doc"""
    +(A::SMat, B::SMat) -> SMat

Return the sum $A + B$.
"""
function +(A::SMat{T}, B::SMat{T}) where T
  nrows(A) != nrows(B) && error("Matrices must have same number of rows")
  ncols(A) != ncols(B) && error("Matrices must have same number of columns")
  C = sparse_matrix(base_ring(A))
  m = min(nrows(A), nrows(B))
  for i=1:m
    push!(C, A[i] + B[i])
  end
  for i=m+1:nrows(A)
    push!(C, A[i])
  end
  for i=m+1:nrows(B)
    push!(C, B[i])
  end
  return C
end

@doc Markdown.doc"""
    -(A::SMat, B::SMat) -> SMat

Return the difference $A - B$.
"""
function -(A::SMat{T}, B::SMat{T}) where T
  nrows(A) != nrows(B) && error("Matrices must have same number of rows")
  ncols(A) != ncols(B) && error("Matrices must have same number of columns")
  C = sparse_matrix(base_ring(A))
  m = min(nrows(A), nrows(B))
  for i=1:m
    push!(C, A[i]-B[i])
  end
  for i=m+1:nrows(A)
    push!(C, A[i])
  end
  if m < nrows(B)
    n = base_ring(B)(-1)
    for i=m+1:nrows(B)
      push!(C, n*B[i])
    end
  end
  return C
end

#TODO: this is a problem in anther ops as well: a length zero row does not have a base_ring
################################################################################
#
#  Scalar multiplication
#
################################################################################

@doc Markdown.doc"""
    *(b::T, A::SMat{T}) -> SMat{T}

Return the product $b \cdot A$.
"""
function *(b::T, A::SMat{T}) where {T <: RingElem}
  B = sparse_matrix(base_ring(A))
  if iszero(b)
    return B
  end
  for a = A
    push!(B, b*a)
  end
  return B
end

@doc Markdown.doc"""
    *(b::Integer, A::SMat{T}) -> SMat{T}

Return the product $b \cdot A$.
"""
function *(b::Integer, A::SMat{T}) where T
  return base_ring(A)(b)*A
end

@doc Markdown.doc"""
    *(b::fmpz, A::SMat{T}) -> SMat{T}

Return the product $b \cdot A$.
"""
function *(b::fmpz, A::SMat{T}) where T
  return base_ring(A)(b)*A
end

# To remove ambiguity we define:
@doc Markdown.doc"""
    *(b::fmpz, A::SMat{fmpz}) -> SMat{fmpz}

Return the product $b \cdot A$.
"""
function *(b::fmpz, A::SMat{fmpz})
  if iszero(b)
    return zero_matrix(SMat, FlintZZ, nrows(A), ncols(A))
  end
  B = sparse_matrix(base_ring(A))
  B.c = ncols(A)
  for a in A
    push!(B, b * a)
  end
  return B
end

################################################################################
#
#  Submatrix
#
################################################################################

@doc Markdown.doc"""
    sub(A::SMat, r::UnitRange, c::UnitRange) -> SMat

Return the submatrix of $A$, where the rows correspond to $r$ and the columns
correspond to $c$.
"""
function sub(A::SMat{T}, r::UnitRange, c::UnitRange) where T
  B = sparse_matrix(base_ring(A))
  B.nnz = 0
  B.c = length(c)
  for i in r
    rw = sparse_row(base_ring(A))
    ra = A.rows[i]
    for j=1:length(ra.values)
      if ra.pos[j] in c
        push!(rw.values, ra.values[j])
        push!(rw.pos, ra.pos[j]-c.start+1)
      end
    end
    push!(B, rw)
  end
  return B
end

#TODO: map to getindex with ranges and slices

################################################################################
#
#  Valence
#
################################################################################

@doc Markdown.doc"""
    valence_mc{T}(A::SMat{T}; extra_prime = 2, trans = Array{SMatSLP_add_row{T}, 1}()) -> T

Uses a Monte-Carlo algorithm to compute the valence of $A$. The valence is the
valence of the minimal polynomial $f$ of $A'*A$, thus the last non-zero
coefficient, typically $f(0)$.

The valence is computed modulo various primes until the computation stabilises
for `extra_prime` many.

`trans`, if given, is  a SLP (straight-line-program) in GL(n, Z). Then the
valence of `trans` * $A$  is computed instead.
"""
function valence_mc(A::SMat{T}; extra_prime = 2, trans = Array{SMatSLP_add_row{T}, 1}()) where T
  # we work in At * A (or A * At) where we choose the smaller of the 2
  # matrices
  if false && ncols(A) > nrows(A)
    At = A
    A = transpose(A)
  else
    At = transpose(A)
  end
  if maximum(abs, A) > 2^20
    mm = mul_mod_big!
    println("mul big case")
  else
    mm = mul_mod!
    println("mul small case")
  end
  c1 = Array{Int}(undef, ncols(A))
  c2 = Array{Int}(undef, nrows(A))

  for i=1:ncols(A)
    c1[i] = Int(rand(-10:10))
  end

  c = copy(c1) ## need the starting value for the other primes as well

  p = next_prime(2^30)
  k = FiniteField(p)
  d = 10
  v = Array{typeof(k(1)), 1}()
  push!(v, k(c1[1]))
  while true
    while length(v) <= d
      mm(c2, A, c1, p)
      if length(trans) >0
        apply_t!(c2, trans)
        apply!(c2, trans)
      end
      mm(c1, At, c2, p)
      push!(v, k(c1[1]))
    end
    f = berlekamp_massey(v)
    fl = true
    if !fl || degree(f) > div(d, 2)-1
      d += 10
      continue
    end
    df = degree(f)
    println("Poly degree is $df, dims $(nrows(A)) x $(ncols(A))")

    V = fmpz(leading_coefficient(f))
    pp = fmpz(p)

    v = Array{typeof(k(1))}(undef, 2*degree(f)+1)
    while true
      p = next_prime(p)
      println(p)
      k = FiniteField(p)
      copy!(c1, c)
      v[1] = k(c1[1])
      for i=1:2*degree(f)
        mm(c2, A, c1, p)
        if length(trans) >0
          apply_t!(c2, trans)
          apply!(c2, trans)
        end
        mm(c1, At, c2, p)
        v[i+1] = k(c1[1])
      end
      f = berlekamp_massey(v)
      fl = true
      if !fl || degree(f) > df
        println("initial guess was wrong...")
        d += 10
        return -1
      end

      Vn = crt(V, pp, fmpz(leading_coefficient(f)), fmpz(p))
      pp *= p
      if 2*Vn > pp
        Vn = Vn - pp
      end
      if Vn == V
        extra_prime -= 1
        if extra_prime == 0
          return V
        end
      end
      V = Vn
    end
  end
end

function valence_mc(A::SMat{T}, p::Int) where T
  # we work in At * A (or A * At) where we choose the smaller of the 2
  # matrices
  if false && ncols(A) > nrows(A)
    At = A
    A = transpose(A)
  else
    At = transpose(A)
  end
  if maximum(abs, A) > 2^20
    mm = mul_mod_big!
    println("mul big case")
  else
    mm = mul_mod!
    println("mul small case")
  end
  c1 = zeros{Int}(ncols(A))
  c2 = zeros{Int}(nrows(A))

  for i=1:ncols(A)
    c1[i] = Int(rand(-10:10))
  end

  c = copy(c1) ## need the starting value for the other primes as well

  k = FiniteField(p)
  d = 10
  v = Array{elem_type(k), 1}()
  push!(v, k(c1[1]))
  while true
    while length(v) <= d
      mm(c2, A, c1, p)
      mm(c1, At, c2, p)
      push!(v, k(c1[1]))
    end
    f = berlekamp_massey(v)
    fl = true
    if !fl || degree(f) > div(d, 2)-1
      d += 10
      continue
    end
    df = degree(f)
    return f
  end
end

################################################################################
#
#  Vertical concatentation
#
################################################################################

@doc Markdown.doc"""
    vcat!(A::SMat, B::SMat) -> SMat

Vertically joins $A$ and $B$ inplace, that is, the rows of $B$ are
appended to $A$.
"""
function vcat!(A::SMat{T}, B::SMat{T}) where T
  @assert length(A.rows) == A.r
  @assert length(B.rows) == B.r
  A.r += B.r
  A.c = max(A.c, B.c)
  A.nnz += B.nnz
  A.rows = vcat(A.rows, B.rows)
  @assert length(A.rows) == A.r
end

@doc Markdown.doc"""
    vcat(A::SMat, B::SMat) -> SMat

Vertically joins $A$ and $B$.
"""
function vcat(A::SMat{T}, B::SMat{T}) where T
  @assert length(A.rows) == A.r
  @assert length(B.rows) == B.r
  C = copy(A)
  vcat!(C, B)
  return C
end

################################################################################
#
#  Horizontal concatentation
#
################################################################################

@doc Markdown.doc"""
    hcat!(A::SMat, B::SMat) -> SMat

Horizontally concatenates $A$ and $B$, inplace, changing $A$.
"""
function hcat!(A::SMat{T}, B::SMat{T}) where T
  o = A.c
  A.c += B.c
  nnz = A.nnz
  for i=1:min(nrows(A), nrows(B))
    for (p,v) in B[i]
      push!(A[i].pos, p+o)
      push!(A[i].values, v)
    end
  end
  for i=min(nrows(A), nrows(B))+1:nrows(B)
    sr = SRow{T}()
    for (p,v) in B[i]
      push!(sr.pos, p+o)
      push!(sr.values, v)
    end
    push!(A, sr)
  end
  A.nnz = nnz + B.nnz #A.nnz may have changed - if B is longer
end

@doc Markdown.doc"""
    hcat(A::SMat, B::SMat) -> SMat

Horizontally concatenates $A$ and $B$.
"""
function hcat(A::SMat{T}, B::SMat{T}) where T
  C = copy(A)
  hcat!(C, B)
  return C
end

################################################################################
#
#  Append sparse row to sparse matrix
#
################################################################################
@doc Markdown.doc"""
    push!(A::SMat{T}, B::SRow{T}) where T

Appends the sparse row ```B``` to ```A```.
"""
function push!(A::SMat{T}, B::SRow{T}) where T
  push!(A.rows, B)
  A.r += 1
  @assert length(A.rows) == A.r
  A.nnz += length(B.pos)
  if length(B.pos) > 0
    A.c = max(A.c, B.pos[end])
  end
end

################################################################################
#
#  Conversion to fmpz_mat
#
################################################################################

@doc Markdown.doc"""
    fmpz_mat{T <: Integer}(A::SMat{T})

The same matrix $A$, but as an `fmpz_mat`.
Requires a conversion from the base ring of $A$ to $\mathbf ZZ$.
"""
function fmpz_mat(A::SMat{T}) where T <: Integer
  B = zero_matrix(FlintZZ, A.r, A.c)
  for i = 1:length(A.rows)
    ra = A.rows[i]
    for j = 1:length(ra.pos)
      B[i, ra.pos[j]] = ra.values[j]
    end
  end
  return B
end

@doc Markdown.doc"""
    fmpz_mat(A::SMat{fmpz})

The same matrix $A$, but as an `fmpz_mat`.
"""
function fmpz_mat(A::SMat{fmpz})
  B = zero_matrix(FlintZZ, A.r, A.c)
  for i = 1:length(A.rows)
    ra = A.rows[i]
    for j = 1:length(ra.pos)
      B[i, ra.pos[j]] = ra.values[j]
    end
  end
  return B
end

################################################################################
#
#  Maximum, minimum and norms
#
################################################################################

@doc Markdown.doc"""
    hadamard_bound2(A::SMat{T}) -> T

The square of the product of the norms of the rows of $A$.
"""
function hadamard_bound2(A::SMat)
  return prod([norm2(x) for x=A])
end

@doc Markdown.doc"""
    maximum(abs, A::SMat{fmpz}) -> fmpz

  Finds the largest, in absolute value, entry of $A$.
"""
function maximum(::typeof(abs), A::SMat{fmpz})
  if length(A.rows) == 0
    return zero(FlintZZ)
  end
  m = abs(A.rows[1].values[1])
  for i in A.rows
    for j in i.values
      if cmpabs(m, j) < 0
        m = j
      end
    end
  end
  return abs(m)
end

@doc Markdown.doc"""
    maximum(A::SMat{T}) -> T

Finds the largest entry of $A$.
"""
function maximum(A::SMat)
  m = zero(base_ring(A))
  for i in A
    for j in i.values
      if cmp(m, j) < 0
        m = j
      end
    end
  end
  return m
end

@doc Markdown.doc"""
    minimum(A::SMat{T}) -> T

Finds the smallest entry of $A$.
"""
function minimum(A::SMat)
  m = zero(base_ring(A))
  for i in A.rows
    for j in i.values
      if cmp(m, j) > 0
        m = j
      end
    end
  end
  return m
end

function maximum(::typeof(nbits), A::SMat{fmpz})
  if length(A.rows) == 0
    return zero(FlintZZ)
  end
  m = nbits(A.rows[1].values[1])
  for i in A.rows
    for j in i.values
      if m < nbits(j)
        m = nbits(j)
      end
    end
  end
  return m
end

################################################################################
#
#  Test upper triangular shape
#
################################################################################

@doc Markdown.doc"""
    isupper_triangular(A::SMat)

Returns true if and only if $A$ is upper (right) triangular.
"""
function isupper_triangular(A::SMat)
  for i=2:A.r
    if iszero(A[i - 1])
      if iszero(A[i])
        continue
      else
        return false
      end
    else
      if iszero(A[i])
        continue
      else
        if A[i - 1].pos[1] >= A[i].pos[1]
          return false
        else
          continue
        end
      end
    end
  end
  return true
end

################################################################################
#
#  Sparse to dense
#
################################################################################

function Nemo.matrix(A::SMat)
  R = base_ring(A)
  M = zero_matrix(R, nrows(A), ncols(A))
  for i=1:nrows(A)
    for j=1:length(A.rows[i].pos)
      M[i, A.rows[i].pos[j]] = A.rows[i].values[j]
    end
  end
  return M
end

################################################################################
#
#  Identity matrices
#
################################################################################

@doc Markdown.doc"""
    identity_matrix(::Type{SMat}, R::Ring, n::Int)

Return a sparse $n$ times $n$ identity matrix over $R$.
"""
function identity_matrix(::Type{SMat}, R::Ring, n::Int)
  A = sparse_matrix(R)
  A.c = n
  for i in 1:n
    push!(A, sparse_row(R, [i], [one(R)]))
  end
  return A
end

function identity_matrix(::Type{MatElem}, R::Ring, n::Int)
  return identity_matrix(R, n)
end

@doc Markdown.doc"""
    zero_matrix(::Type{SMat}, R::Ring, n::Int)

Return a sparse $n$ times $n$ zero matrix over $R$.
"""
function zero_matrix(::Type{SMat}, R::Ring, n::Int)
  S = sparse_matrix(R)
  S.rows = [sparse_row(R) for i=1:n]
  S.c = S.r = n
  return S
end

@doc Markdown.doc"""
    zero_matrix(::Type{SMat}, R::Ring, n::Int, m::Int)

Return a sparse $n$ times $m$ zero matrix over $R$.
"""
function zero_matrix(::Type{SMat}, R::Ring, n::Int, m::Int)
  S = sparse_matrix(R)
  S.rows = [sparse_row(R) for i=1:n]
  S.r = n
  S.c = m
  return S
end

################################################################################
#
#  Julias concatenatin syntax
#
################################################################################

function Base.cat(n::Int, A::SMat...)
  if n == 1
    return vcat(A...)
  elseif n == 2
    return hcat(A...)
  else
    error("dims must be 1 or 2")
  end
end

function Base.vcat(A::SMat...)
  B = copy(A[1])
  for i=2:length(A)
    for r = A[i].rows
      push!(B, copy(r))
    end
  end
  return B
end

#base case
function Base.hcat(A::SMat)
  return A
end

function Base.hcat(A::SMat, B::SMat, C::SMat...)
  return hcat(hcat(A, B), C...)
end

function Base.cat(dims::Tuple{Int, Int}, A::SMat...)
  if dims != (1,2)
    error("dims must be (1, 2)")
  end
  B = copy(A[1])
  c = B.c
  for i=2:length(A)
    for j=1:nrows(A[i])
      R = SRow(base_ring(B))
      for (k,v) = A[i].rows[j]
        push!(R.pos, k+c)
        push!(R.values, v)
      end
      push!(B, R)
    end
    c += A[i].c
  end
  return B
end

################################################################################
#
#  Is one/zero?
#
################################################################################

@doc Markdown.doc"""
    isone(A::SMat)

Tests if $A$ is an identity matrix.
"""
function isone(A::SMat)
  if ncols(A) != nrows(A)
    return false
  end
  for (i, r) in enumerate(A)
    if length(r.pos) != 1
      return false
    end
    if r.pos[1] != i || !isone(r.values[1])
      return false
    end
  end
  return true
end

@doc Markdown.doc"""
    iszero(A::SMat)

Tests if $A$ is a zero matrix.
"""
function iszero(A::SMat)
  for r in A
    if !iszero(r)
      return false
    end
  end
  return true
end

################################################################################
#
#  File Serialization
#
################################################################################

@doc Markdown.doc"""
    toNemo(io::IOStream, A::SMat; name = "A")

  Prints the SMat as a julia-program into the file corresponding to `io`.
  The file can be included to get the matrix.
  `name` controls the variable name of the matrix.
"""
function toNemo(io::IOStream, A::SMat; name = "A")
  T = typeof(A.rows[1].values[1])
  println(io, name, " = SMat{$T}()")
  for i=A.rows
    print(io, "push!($(name), SRow{$T}(Tuple{Int, $T}[");
    for j=1:length(i.pos)-1
      print(io, "($(i.pos[j]), $(i.values[j])), ")
    end
    println(io, "($(i.pos[end]), $(i.values[end]))]))")
  end
end

@doc Markdown.doc"""
    toNemo(io::String, A::SMat; name = "A")

  Prints the SMat as a julia-program into the file named `io`.
  The file can be included to get the matrix.
  `name` controls the variable name of the matrix.
"""
function toNemo(f::String, A::SMat; name = "A")
  io = open(f, "w")
  toNemo(io, A, name=name)
  close(io)
end

################################################################################
#
#  Conversion to julia types
#
################################################################################

@doc Markdown.doc"""
    sparse(A::SMat) -> SparseMatrixCSC

The same matrix, but as a sparse matrix of julia type `SparseMatrixCSC`.
"""
function SparseArrays.sparse(A::SMat{T}) where T
  I = zeros(Int, A.nnz)
  J = zeros(Int, A.nnz)
  V = Vector{T}(undef, A.nnz)
  i = 1
  for r = 1:nrows(A)
    for j=1:length(A.rows[r].pos)
      I[i] = r
      J[i] = A.rows[r].pos[j]
      V[i] = A.rows[r].values[j]
      i += 1
    end
  end
  return SparseArrays.sparse(I, J, V)
end

@doc Markdown.doc"""
    Array(A::SMat{T}) -> Array{T, 2}

The same matrix, but as a two-dimensional julia array.
"""
function Array(A::SMat{T}) where T
  R = zero_matrix(base_ring(A), A.r, A.c)
  for i=1:nrows(A)
    for j=1:length(A.rows[i].pos)
      R[i, A.rows[i].pos[j]] = A.rows[i].values[j]
    end
  end
  return R
end

#TODO: write a kronnecker-row-product, this is THE
# Kronecker product corresponding to the matrix of a tensor
# product of homs represented as column vectors
#in Oscar, we should be doing row vectors...
function kronecker_product(A::SMat{T}, B::SMat{T}) where {T}
  C = sparse_matrix(base_ring(A))

  for rA = A.rows
    for rB = B.rows
      p = Int[]
      v = T[]
      o = (rA.pos[1]-1)*ncols(B)
      for (pp, vv) = rA
        for (qq, ww) = rB
          push!(p, qq+o)
          push!(v, vv*ww)
        end
        o += ncols(B)
      end
      push!(C, sparse_row(base_ring(A), p, v))
    end
  end
  return C
end


text:
# Import the prata package
using Pkg
Pkg.activate("Build/Julia/prata.jl")
using prata

# Connect to a host, this could be run running or one you plan on running
println("TEST: Basic pulbish/subscribe")
m = prata.connect("127.0.0.1",25565)

# This is used to launch a host
prata.host(m)

`
PUBLISHER
`

# Create a publisher
publisher = prata.publisher(m)

# publish to a channel, if this channel does not exist it will be created
prata.publish(publisher, "My Channel", "Information")

`
SUBSCRIBER
`

# Create a subscriber
subscriber = prata.subscriber(m)

# Listen from the channel
data = prata.listen(subscriber, "My Channel")

if data != "Information"
    println("Failure in Julia. Got: ", data, " Needed Information")
end

# Terminate the entire prata network, stopping all channels.
prata.terminate(m)

max_stars_repo_name: roboscienceorg/Prata

max_issues_repo_name: roboscienceorg/Prata

max_forks_repo_name: roboscienceorg/Prata


text:
# check if package loads
using PhotoOrganizer


text:
module QuOptimalControl

include("problems.jl")
include("cost_functions.jl")
include("tools.jl")
include("grape_tools.jl")
include("timeevolution.jl")
include("algorithms.jl")
include("dCRAB.jl")
include("GRAPE.jl")
include("solve.jl")
include("visualisation.jl")

export C1, C2, C3, C4, C5, C6, C7, fom_func
export init_GRAPE, init_ensemble
export commutator, eig_factors, expm_exact_gradient, trace_matmul, save, load
export pw_evolve,
    pw_evolve_save,
    pw_ham_save!,
    pw_gen_save!,
    pw_prop_save!,
    evolve_func,
    evolve_func!
export IntegratorType, Piecewise, Continuous
export algorithm,
    gradientBased, gradientFree, SolutionResult, GRAPE, ADGRAPE, EnsembleSolutionResult
export solve
export grad_func
export Problem,
    SystemType,
    ClosedSystem,
    OpenSystem,
    StateTransfer,
    UnitaryGate,
    CoherenceTransfer,
    EnsembleProblem,
    Experiment,
    ExperimentInterface
export visualise_expt_val, visualise_pulse

end


text:
#######################################################################
#
# An example of creating an Excel chart in a chartsheet with XlsxWriter.
#
# Original Python Copyright 2013-2016, John McNamara, jmcnamara@cpan.org
# https://github.com/jmcnamara/XlsxWriter

using XlsxWriter

function test()

	wb = Workbook("chartsheet.xlsx")
	ws = add_worksheet!(wb)
	bold = add_format!(wb, Dict("bold"=>1))

	# Add a chartsheet. A worksheet that only holds a chart.
	cs = add_chartsheet!(wb)

	# Add the worksheet data that the charts will refer to.
	headings = ["Number" "Batch 1" "Batch 2"]
	data = [
		[2 3 4 5 6 7];
		[10 40 50 20 10 50];
		[30 60 70 50 40 30];
	]

	write_row!(ws, "A1", headings, fmt=bold)
	write_column!(ws, "A2", data[1,:])
	write_column!(ws, "B2", data[2,:])
	write_column!(ws, "C2", data[3,:])

	# Create a new bar chart.
	chart1 = add_chart!(wb, Dict("type"=> "bar"))

	# Configure the first series.
	add_series!(chart1, Dict(
		"name"=>       "=Sheet1!\$B\$1",
		"categories"=> "=Sheet1!\$A\$2:\$A\$7",
		"values"=>     "=Sheet1!\$B\$2:\$B\$7",
	))

	# Configure a second series. Note use of alternative syntax to define ranges.
	add_series!(chart1, Dict(
		"name"=>       ["Sheet1", 0, 2],
		"categories"=> ["Sheet1", 1, 0, 6, 0],
		"values"=>     ["Sheet1", 1, 2, 6, 2],
	))

	# Add a chart title and some axis labels.
	set_title!(chart1, Dict("name"=> "Results of sample analysis"))
	set_x_axis!(chart1, Dict("name"=> "Test number"))
	set_y_axis!(chart1, Dict("name"=> "Sample length (mm)"))

	# Set an Excel chart style.
	set_style!(chart1, 11)

	# Add the chart to the chartsheet.
	set_chart!(cs, chart1)

	close(wb)
	isfile("chartsheet.xlsx")
end

test()


text:
__precompile__()

module POMDPToolbox

using POMDPs

import POMDPs: Updater, update, initialize_belief, pdf, mode, updater
import POMDPs: Simulator, simulate
import POMDPs: action, value, solve
import POMDPs: actions, action_index, state_index, obs_index, iterator, sampletype, states, n_actions, n_states, observations, n_observations, discount, isterminal
import POMDPs: generate_sr, initial_state
import POMDPs: implemented
import Base: rand, rand!, mean, ==

using ProgressMeter
using StatsBase
using DataFrames

# export commons
export
    update,
    initialize_belief,
    create_belief,
    pdf,
    updater,
    simulate,
    action,
    solve

# only include the things that are working
# export things immediately above the file they are contained in

# beliefs
export
    VoidUpdater
include("beliefs/void.jl")

export
    DiscreteBelief,
    DiscreteUpdater,
    uniform_belief,
    product
include("beliefs/discrete.jl")

export
    PreviousObservationUpdater,
    FastPreviousObservationUpdater,
    PrimedPreviousObservationUpdater
include("beliefs/previous_observation.jl")

export
    KMarkovUpdater

include("beliefs/k_previous_observations.jl")

export
    Particle,
    ParticleBelief,
    ParticleDistribution,
    SIRParticleUpdater,
    mode
include("beliefs/particle.jl")

# convenience
include("convenience/implementations.jl")

# policies
export
    AlphaVectorPolicy,
    belief_vector,
    unnormalized_util
include("policies/alpha_vector.jl")

export
    VectorPolicy,
    VectorSolver,
    ValuePolicy
include("policies/vector.jl")

export
    RandomPolicy,
    RandomSolver
include("policies/random.jl")

export
    StochasticPolicy,
    UniformRandomPolicy,
    CategoricalTabularPolicy,
    EpsGreedyPolicy
include("policies/stochastic.jl")

export
    FunctionPolicy,
    FunctionSolver
include("policies/function.jl")

export
    PolicyWrapper,
    payload
include("policies/utility_wrapper.jl")

# simulators
export RolloutSimulator
include("simulators/rollout.jl")

export
    SimHistory,
    POMDPHistory,
    MDPHistory,
    AbstractPOMDPHistory,
    AbstractMDPHistory,
    HistoryIterator,
    eachstep,
    state_hist,
    action_hist,
    observation_hist,
    belief_hist,
    reward_hist,
    info_hist,
    ainfo_hist,
    uinfo_hist,
    exception,
    backtrace,
    undiscounted_reward,
    discounted_reward,
    n_steps,
    step_tuple
include("simulators/history.jl")

export sim
include("simulators/sim.jl")

export HistoryRecorder
include("simulators/history_recorder.jl")

export
    StepSimulator,
    stepthrough
include("simulators/stepthrough.jl")

export
    Sim,
    run,
    run_parallel,
    problem
include("simulators/parallel.jl")

# model tools
export
    generate_sri,
    generate_sori,
    action_info,
    solve_info,
    update_info
include("model/info.jl")

export uniform_state_distribution
include("model/initial.jl")

export
    ordered_states,
    ordered_actions,
    ordered_observations
include("model/ordered_spaces.jl")

export GenerativeBeliefMDP
include("model/generative_belief_mdp.jl")

export FullyObservablePOMDP
include("model/fully_observable_pomdp.jl")

export UnderlyingMDP
include("model/underlying_mdp.jl")

# tools for distributions
include("distributions/distributions_jl.jl")

export obs_weight
include("model/obs_weight.jl")

export
    weighted_iterator
include("distributions/weighted_iteration.jl")

export
    SparseCat
include("distributions/sparse_cat.jl")

export
    BoolDistribution
include("distributions/bool.jl")

# testing
export test_solver
include("testing/solver.jl")

export
    probability_check,
    obs_prob_consistency_check,
    trans_prob_consistency_check
include("testing/model.jl")

end # module


text:
struct WeakRefShim{T}
    ref::WeakRef
    WeakRefShim(obj::T) where {T} = new{T}(WeakRef(obj))
end

valueof(shim::WeakRefShim) = getfield(shim, :ref).value

@noinline error_finalized() =
    error("object to be finalized cannot be accessed in async finalizer")

"""
    AsyncFinalizers.unsafe_unwrap(shim) -> object

Unwrap a `shim` and obtain the original object. The user is responsible for ensuring that
`object` does not escape.

See [`AsyncFinalizers.onfinalize`](@ref)
"""
function AsyncFinalizers.unsafe_unwrap(shim::WeakRefShim{T}) where {T}
    obj = valueof(shim)
    if obj isa T
        return obj
    else
        error_finalized()
    end
end

function Base.getproperty(shim::WeakRefShim{T}, name::Symbol) where {T}
    obj = valueof(shim)
    if obj isa T
        return getproperty(obj, name)
    else
        error_finalized()
    end
end

function Base.getproperty(shim::WeakRefShim{T}, name::Symbol, order::Symbol) where {T}
    obj = valueof(shim)
    if obj isa T
        return getproperty(obj, name, order)
    else
        error_finalized()
    end
end

function Base.setproperty!(shim::WeakRefShim{T}, name::Symbol, x) where {T}
    obj = valueof(shim)
    if obj isa T
        return setproperty!(obj, name, x)
    else
        error_finalized()
    end
end

function Base.setproperty!(shim::WeakRefShim{T}, name::Symbol, x, order::Symbol) where {T}
    obj = valueof(shim)
    if obj isa T
        return setproperty!(obj, name, x, order)
    else
        error_finalized()
    end
end


text:
module TestUtils

using Test, DataFrames

@testset "make_unique" begin
    @test DataFrames.make_unique([:x, :x, :x_1, :x2], makeunique=true) == [:x, :x_2, :x_1, :x2]
    @test_throws ArgumentError DataFrames.make_unique([:x, :x, :x_1, :x2], makeunique=false)
    @test DataFrames.make_unique([:x, :x_1, :x2], makeunique=false) == [:x, :x_1, :x2]
end

@testset "repeat count" begin
    df = DataFrame(a = 1:2, b = 3:4)
    ref = DataFrame(a = repeat(1:2, 2),
                    b = repeat(3:4, 2))
    @test repeat(df, 2) == ref
    @test repeat(view(df, 1:2, :), 2) == ref

    @test size(repeat(df, 0)) == (0, 2)
    @test size(repeat(df, false)) == (0, 2)
    @test_throws ArgumentError repeat(df, -1)
end

@testset "repeat inner_outer" begin
    df = DataFrame(a = 1:2, b = 3:4)
    ref = DataFrame(a = repeat(1:2, inner = 2, outer = 3),
                    b = repeat(3:4, inner = 2, outer = 3))
    @test repeat(df, inner = 2, outer = 3) == ref
    @test repeat(view(df, 1:2, :), inner = 2, outer = 3) == ref

    @test size(repeat(df, inner = 2, outer = 0)) == (0, 2)
    @test size(repeat(df, inner = 0, outer = 3)) == (0, 2)
    @test size(repeat(df, inner = 2, outer = false)) == (0, 2)
    @test size(repeat(df, inner = false, outer = 3)) == (0, 2)

    @test_throws ArgumentError repeat(df, inner = 2, outer = -1)
    @test_throws ArgumentError repeat(df, inner = -1, outer = 3)
end

@testset "repeat! count" begin
    df = DataFrame(a = 1:2, b = 3:4)
    ref = DataFrame(a = repeat(1:2, 2),
                    b = repeat(3:4, 2))
    a = df.a
    b = df.b
    repeat!(df, 2)
    @test df == ref
    @test a == 1:2
    @test b == 3:4

    for v in (0, false)
        df = DataFrame(a = 1:2, b = 3:4)
        repeat!(df, v)
        @test size(df) == (0, 2)
    end

    df = DataFrame(a = 1:2, b = 3:4)
    @test_throws ArgumentError repeat(df, -1)
    @test df == DataFrame(a = 1:2, b = 3:4)

    @test_throws MethodError repeat!(view(df, 1:2, :), 2)
end

@testset "repeat! inner_outer" begin
    df = DataFrame(a = 1:2, b = 3:4)
    ref = DataFrame(a = repeat(1:2, inner = 2, outer = 3),
                    b = repeat(3:4, inner = 2, outer = 3))
    a = df.a
    b = df.b
    repeat!(df, inner = 2, outer = 3)
    @test df == ref
    @test a == 1:2
    @test b == 3:4

    for v in (0, false)
        df = DataFrame(a = 1:2, b = 3:4)
        repeat!(df, inner = 2, outer = v)
        @test size(df) == (0, 2)

        df = DataFrame(a = 1:2, b = 3:4)
        repeat!(df, inner = v, outer = 3)
        @test size(df) == (0, 2)
    end

    df = DataFrame(a = 1:2, b = 3:4)
    @test_throws ArgumentError repeat(df, inner = 2, outer = -1)
    @test_throws ArgumentError repeat(df, inner = -1, outer = 3)
    @test df == DataFrame(a = 1:2, b = 3:4)

    @test_throws MethodError repeat!(view(df, 1:2, :), inner = 2, outer = 3)
end

@testset "funname" begin
    @test DataFrames.funname(sum ∘ skipmissing ∘ Base.div12) ==
          :sum_skipmissing_div12
end

@testset "pre-Julia 1.3 @spawn replacement" begin
    t = @sync DataFrames.@spawn begin
        sleep(1)
        true
    end
    @test fetch(t) === true
end

end # module


text:
# genGauss
@testset "genGauss" begin
    # One-dimensional test
    aG = ones(3);
    mG = zeros(1,3);
    mG[1,:] = [0.0, 0.7, -0.3];
    AG = zeros(1,1,3);
    AG[1,1,:] = [100 80 120];
    g = FPRS.genGauss(aG, mG, AG);
    @test abs.(sum(g) - 2.8286690853954575) < 67.0 * eps()

    # Two-dimensional test
    # Generate 3 Gaussians
    aG = ones(3);
    mG = zeros(2,3);
    mG[1,:] = [0.0, 0.7, -0.3];
    mG[2,:] = [0.0, -0.6, -0.06];
    AG = zeros(2,2,3);
    AG[:,:,1] = [100 0.05; 0 100];
    AG[:,:,2] = [100 0.15; 0 100];
    AG[:,:,3] = [100 -0.3; 0.4 100];
    g = FPRS.genGauss(aG, mG, AG);
    @test abs.(sum(g) - 102.63572536664947) < 67.0^2 * eps()

    xmin = [-1.0 -2.0];
    xmax = [2.0   1.0];
    nPts = [127 127];
    g = FPRS.genGauss(aG, mG, AG, nPts = nPts, xmin=xmin, xmax=xmax);
    @test abs.(sum(g) - 166.25310747322118) < 127.0^2 * eps()

    # Test errors
    aG = ones(2)
    @test_throws DimensionMismatch FPRS.genGauss(aG, mG, AG);
    aG = ones(3)
    AG = zeros(2,3,3)
    @test_throws DimensionMismatch FPRS.genGauss(aG, mG, AG);

    # Three-dimensional test
    aG = ones(3);
    mG = zeros(3,3);
    mG[1,:] = [0.0, 0.7, -0.3];
    mG[2,:] = [0.0, -0.6, -0.06];
    mG[3,:] = [0.3, -0.5, -0.2];
    AG = zeros(3,3,3);
    AG[:,:,1] = [100 0.05 50; 0   100 20; 0  0  100];
    AG[:,:,2] = [100 0.15 10; 0   100 10; 0  20 100];
    AG[:,:,3] = [100 -0.3 10; 0.4 100 50; 50 0  100];
    g = FPRS.genGauss(aG, mG, AG);
    @test abs.(sum(g) - 627.7968551738583) < 67.0^3 * eps()
end

# genRandGauss
@testset "genRandGauss" begin
    # 1D
    g = FPRS.genRandGauss(nGauss = 50, nPts=[512], AGmin=100, AGmax=1000, mGrad=0.5, randSeed=0, suppShape="box")
    @test abs.(sum(g) - 63.77852000489939) < 512 * eps()

    # 2D
    g = FPRS.genRandGauss(nGauss = 10, nPts=[51,51], mGrad=0.5, randSeed=0, suppShape="ball")
    @test abs.(sum(g) - 45.34679104451378) < 51^2 * eps()

    # 3D
    g = FPRS.genRandGauss(nGauss = 10, nPts=[31,31,21], randSeed=0, suppShape="ball")
    @test abs.(sum(g) - 17.801826001001068) < 31^2 * 21 * eps()

    # Reverse-engineered. (?) Why not ArgumentError
    @test_throws MethodError FPRS.genGauss(suppShape = 1);
end

max_stars_repo_name: tsipenyuk/FPRS

max_issues_repo_name: tsipenyuk/FPRS

max_forks_repo_name: tsipenyuk/FPRS