text:
"""
`module  JAC.ReducedDensityMatrix`
    ... a submodel of JAC that contains all methods for computing reduced density matrices, natural orbitals, radial density
        distributions and related information for some level(s).
"""
module ReducedDensityMatrix

    using Printf, ..Basics, ..Defaults, ..InteractionStrength, ..ManyElectron, ..Nuclear, ..Radial, ..RadialIntegrals,
                  ..SpinAngular, ..TableStrings

    """
    `struct  ReducedDensityMatrix.Settings  <:  AbstractPropertySettings`
        ... defines a type for providing and parameters for reduced density matrices and other parameters.

        + calc1pRDM                ::Bool             ... True if 1-particle RDM need to be calculated, and false otherwise.
        + calc2pRDM                ::Bool             ... True if 2-particle RDM need to be calculated, and false otherwise.
        + calcCumulant             ::Bool             ... True if the cumulant need to be calculated, and false otherwise.
        + calcNatural              ::Bool             ... True if natural orbitals need to be calculated, and false otherwise.
        + calcDensity              ::Bool             ... True if the radial density need to be calculated, and false otherwise.
        + calcIpq                  ::Bool             ... True if orbital interaction I_pq need to be calculated, and false otherwise.
        + printBefore              ::Bool             ... True if a list of selected levels is printed before the actual computations start.
        + levelSelection           ::LevelSelection   ... Specifies the selected levels, if any.
    """
    struct Settings  <:  AbstractPropertySettings
        calc1pRDM                  ::Bool
        calc2pRDM                  ::Bool
        calcCumulant               ::Bool
        calcNatural                ::Bool
        calcDensity                ::Bool
        calcIpq                    ::Bool
        printBefore                ::Bool
        levelSelection             ::LevelSelection
    end

    """
    `ReducedDensityMatrix.Settings(; calc1pRDM::Bool=true,` calc2pRDM::Bool=false, calcCumulant::Bool=false, calcNatural::Bool=false,
                                     calcDensity::Bool=false, calcIpq::Bool=false, printBefore::Bool=true,
                                     levelSelection::LevelSelection=LevelSelection())
        ... keyword constructor to overwrite selected value of reduced-density matrix computations.
    """
    function Settings(; calc1pRDM::Bool=true, calc2pRDM::Bool=false, calcCumulant::Bool=false, calcNatural::Bool=false,
                        calcDensity::Bool=false, calcIpq::Bool=false, printBefore::Bool=true,
                        levelSelection::LevelSelection=LevelSelection())
        Settings(calc1pRDM, calc2pRDM, calcCumulant, calcNatural, calcDensity, calcIpq, printBefore, levelSelection)
    end

    # `Base.show(io::IO, settings::ReducedDensityMatrix.Settings)`
    #       ... prepares a proper printout of the variable settings::ReducedDensityMatrix.Settings.
    function Base.show(io::IO, settings::ReducedDensityMatrix.Settings)
        println(io, "calc1pRDM:                $(settings.calc1pRDM)  ")
        println(io, "calc2pRDM:                $(settings.calc2pRDM)  ")
        println(io, "calcCumulant:             $(settings.calcCumulant)  ")
        println(io, "calcNatural:              $(settings.calcNatural)  ")
        println(io, "calcDensity:              $(settings.calcDensity)  ")
        println(io, "calcIpq:                  $(settings.calcIpq)  ")
        println(io, "printBefore:              $(settings.printBefore)  ")
        println(io, "levelSelection:           $(settings.levelSelection)  ")
    end

    """
    `struct  ReducedDensityMatrix.Outcome`
        ... defines a type to keep the outcome of a reduced density-matrix and natural-orbital computation.

        + level                     ::Level                 ... Atomic level to which the outcome refers to.
        + reduced1p                 ::Array{Float64,2}      ... 1-particle RDM (ip,iq)
        + reduced2p                 ::Array{Float64,4}      ... 1-particle RDM (ip,iq,ir,is)
        + orbitalInteraction        ::Array{Float64,2}      ... orbital interaction I_pq = (ip,iq)
        + naturalOccupation         ::Array{Float64,1}      ... Occupation numbers of natural orbitals.
        + naturalSubshells          ::Array{Subshell,1}     ... List of natural orbitals (subshells).
        + naturalOrbitals           ::Dict{Subshell, Orbital} . Dictionary of (one-electron) natural orbitals.
        + electronDensity           ::Radial.Density        ... Radial density distribution.
    """
    struct Outcome
        level                       ::Level
        reduced1p                   ::Array{Float64,2}
        reduced2p                   ::Array{Float64,4}
        orbitalInteraction          ::Array{Float64,2}
        naturalOccupation           ::Array{Float64,1}
        naturalSubshells            ::Array{Subshell,1}
        naturalOrbitals             ::Dict{Subshell, Orbital}
        electronDensity             ::Radial.Density
    end

    """
    `ReducedDensityMatrix.Outcome()`  ... constructor for an `empty` instance of Hfs.Outcome for the computation of isotope-shift properties.
    """
    function Outcome()
        Outcome(Level(), zeros(1,1), zeros(1,1,1,1), zeros(1,1), Float64[], Subshell[], Dict{Subshell, Orbital}(), Radial.Density() )
    end

    # `Base.show(io::IO, outcome::ReducedDensityMatrix.Outcome)`  ... prepares a proper printout of the variable outcome::ReducedDensityMatrix.Outcome.
    function Base.show(io::IO, outcome::ReducedDensityMatrix.Outcome)
        println(io, "level:                   $(outcome.level)  ")
        println(io, "reduced1p:               $(outcome.reduced1p)  ")
        println(io, "reduced2p:               $(outcome.reduced2p)  ")
        println(io, "orbitalInteraction:      $(outcome.orbitalInteraction)  ")
        println(io, "naturalOccupation:       $(outcome.naturalOccupation)  ")
        println(io, "naturalSubshells:        $(outcome.naturalSubshells)  ")
        println(io, "naturalOrbitals:         $(outcome.naturalOrbitals)  ")
        println(io, "electronDensity:         $(outcome.electronDensity)  ")
    end

    """
    `ReducedDensityMatrix.computeProperties(outcome::ReducedDensityMatrix.Outcome, nm::Nuclear.Model, grid::Radial.Grid,
                                            settings::ReducedDensityMatrix.Settings)
        ... to compute all properties for a given level; an outcome::ReducedDensityMatrix.Outcome is returned for
            which the properties are now evaluated explicitly.
    """
    function  computeProperties(outcome::ReducedDensityMatrix.Outcome, nm::Nuclear.Model, grid::Radial.Grid,
                                settings::ReducedDensityMatrix.Settings)
        reduced1p = outcome.reduced1p;     naturalOccupation = outcome.naturalOccupation;
        reduced2p = outcome.reduced2p;     naturalSubshells  = outcome.naturalSubshells;     naturalOrbitals = outcome.naturalOrbitals;
        orbitalInteraction = outcome.orbitalInteraction;       electronDensity = outcome.electronDensity;

        if  settings.calc1pRDM      @warn("Option calc1pRDM not yet implemented.")
        end

        if  settings.calc2pRDM      @warn("Option calc2pRDM not yet implemented.")
        end

        if  settings.calcCumulant   println("Option calcCumulant not yet implemented.")
        end

        if  settings.calcNatural    println("Option calcNatural not yet implemented.")
        end

        if  settings.calcDensity    println("Option calcDensity not yet implemented.")
        end

        if  settings.calcIpq        println("Option calcIpq not yet implemented.")
        end

        newOutcome = ReducedDensityMatrix.Outcome( outcome.level, reduced1p, reduced2p, orbitalInteraction,
                                                   naturalOccupation, naturalSubshells, naturalOrbitals, electronDensity )
        return( newOutcome )
    end

    """
    `ReducedDensityMatrix.computeOutcomes(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                                          settings::ReducedDensityMatrix.Settings; output=true)`
        ... to compute (as selected) the 1-particle and 2-particle reduced density matrices, natural orbitals or other
            requested information for the levels of interest.
    """
    function computeOutcomes(multiplet::Multiplet, nm::Nuclear.Model, grid::Radial.Grid,
                             settings::ReducedDensityMatrix.Settings; output=true)
        println("")
        printstyled("ReducedDensityMatrix.computeOutcomes(): The computation of the reduced density matrices starts now ... \n", color=:light_green)
        printstyled("------------------------------------------------------------------------------------------------------ \n", color=:light_green)
        #
        outcomes = ReducedDensityMatrix.determineOutcomes(multiplet, settings)
        # Display all selected levels before the computations start
        if  settings.printBefore    ReducedDensityMatrix.displayOutcomes(outcomes, settings)    end
        # Calculate all requested properties
        newOutcomes = ReducedDensityMatrix.Outcome[]
        for  outcome in outcomes
            newOutcome = ReducedDensityMatrix.computeProperties(outcome, nm, grid, settings)
            push!( newOutcomes, newOutcome)
        end
        # Print all results to screen
        ReducedDensityMatrix.displayResults(stdout, newOutcomes, settings)
        printSummary, iostream = Defaults.getDefaults("summary flag/stream")
        if  printSummary    ReducedDensityMatrix.displayResults(iostream, newOutcomes, settings)   end
        #
        if    output    return( newOutcomes )
        else            return( nothing )
        end
    end

    """
    `ReducedDensityMatrix.determineOutcomes(multiplet::Multiplet, settings::ReducedDensityMatrix.Settings)`
        ... to determine a list of Outcomes's for the computation of reduced density matrices and other information
            for the given multiplet. It takes into account the particular selections and settings.
            An Array{ReducedDensityMatrix.Outcome,1} is returned. Apart from the level specification, all physical
            properties are set to zero during the initialization process.
    """
    function  determineOutcomes(multiplet::Multiplet, settings::ReducedDensityMatrix.Settings)
        outcomes = ReducedDensityMatrix.Outcome[]
        @show settings.levelSelection
        for  level  in  multiplet.levels
            if  Basics.selectLevel(level, settings.levelSelection)
                push!( outcomes, ReducedDensityMatrix.Outcome(level, zeros(1,1), zeros(1,1,1,1), zeros(1,1),
                                                              Float64[], Subshell[], Dict{Subshell, Orbital}(), Radial.Density()) )
            end
        end
        return( outcomes )
    end

    """
    `ReducedDensityMatrix.displayOutcomes(outcomes::Array{ReducedDensityMatrix.Outcome,1}, settings::ReducedDensityMatrix.Settings)`
        ... to display a list of levels that have been selected for the computations. A small neat table of all
            selected levels and their energies is printed but nothing is returned otherwise. Moreover, the
            selected properties and information is printed as well, though not yet calculated.
    """
    function  displayOutcomes(outcomes::Array{ReducedDensityMatrix.Outcome,1}, settings::ReducedDensityMatrix.Settings)
        nx = 43
        println(" ")
        println("  Selected levels for reduced density matrix & natural orbital calculations:")
        println(" ")
        println("  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  "
        sa = sa * TableStrings.center(10, "Level"; na=2);                             sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center(10, "J^P";   na=4);                             sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(14, "Energy"; na=4);
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=4)
        println(sa);    println(sb);    println("  ", TableStrings.hLine(nx))
        #
        for  outcome in outcomes
            sa  = "  ";    sym = Basics.LevelSymmetry( outcome.level.J, outcome.level.parity)
            sa = sa * TableStrings.center(10, TableStrings.level(outcome.level.index); na=2)
            sa = sa * TableStrings.center(10, string(sym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", outcome.level.energy)) * "    "
            println( sa )
        end
        println("  ", TableStrings.hLine(nx))
        println("\n  The following properties and information will be computed:\n")
        if  settings.calc1pRDM      println("  ++ compute 1-particle RDM.")                       end
        if  settings.calc2pRDM      println("  ++ compute 2-particle RDM.")                       end
        if  settings.calcCumulant   println("  ++ compute cumulant.")                             end
        if  settings.calcNatural    println("  ++ compute natural orbitals & occupation.")        end
        if  settings.calcDensity    println("  ++ compute radial density.")                       end
        if  settings.calcIpq        println("  ++ compute orbital interaction I(p,q).")           end
        #
        return( nothing )
    end

    """
    `ReducedDensityMatrix.displayResults(stream::IO, outcomes::Array{ReducedDensityMatrix.Outcome,1},
                                         settings::ReducedDensityMatrix.Settings)`
        ... to display the reduced density matrices, natural orbitals, etc. for the selected levels.
            A neat table is printed but nothing is returned otherwise.
    """
    function  displayResults(stream::IO, outcomes::Array{ReducedDensityMatrix.Outcome,1}, settings::ReducedDensityMatrix.Settings)
        nx = 133
        println(stream, " ")
        println(stream, "  ReducedDensityMatrix parameters and information:  !! Not yet worked out !! ")
        println(stream, " ")
        println(stream, "  ", TableStrings.hLine(nx))
        sa = "  ";   sb = "  ";   sc = "  "
        sa = sa * TableStrings.center(10, "Level"; na=2);                             sb = sb * TableStrings.hBlank(12)
        sa = sa * TableStrings.center(10, "J^P";   na=4);                             sb = sb * TableStrings.hBlank(14)
        sa = sa * TableStrings.center(14, "Energy"; na=5)
        sb = sb * TableStrings.center(14, TableStrings.inUnits("energy"); na=5);  sc = sc * TableStrings.hBlank(45)
        sa = sa * TableStrings.center(11, "K_nms"; na=4)
        sb = sb * TableStrings.center(11, "[GHz u]" ; na=4)
        sc = sc * TableStrings.center(11, "[a.u.]" ; na=4)
        println(stream, sa);    println(stream, sb);    ## println(stream, sc);
        println(stream, "  ", TableStrings.hLine(nx))
        #
        for  outcome in outcomes
            sa  = "  ";    sym = LevelSymmetry( outcome.level.J, outcome.level.parity)
            sa = sa * TableStrings.center(10, TableStrings.level(outcome.level.index); na=2)
            sa = sa * TableStrings.center(10, string(sym); na=4)
            sa = sa * @sprintf("%.8e", Defaults.convertUnits("energy: from atomic", outcome.level.energy))            * "    "
            println(stream, sa )
            ## sb = TableStrings.hBlank(47)
        end
        println(stream, "  ", TableStrings.hLine(nx), "\n\n")

        return( nothing )
    end

end # module


text:
# test just package loading as of now
using JuliaRunClient
#using Test

# write your own tests here
#@test 1 == 2

ext: jl


text:
export UnitCell
export make_unitcell, makeunitcell
export dimension

export numsite, sitecount,
       hassite,
       addsite!,
       getsite,
       getsiteindex,
       getsitecoord,
       getsiteindexcoord,
       getsitename

export carte2fract,
       fract2carte,
       whichunitcell,
       momentumgrid

export findsiteindex

import LinearAlgebra

export numorbital,
       hasorbital,
       addorbital!,
       getorbital,
       getorbitalindex,
       getorbitalcoord,
       getorbitalindexcoord,
       getorbitalname

@deprecate numorbital(args...) numsite(args...)
@deprecate hasorbital(args...) hassite(args...)
@deprecate addorbital!(args...) addsite!(args...)
@deprecate getorbital(args...) getsite(args...)
@deprecate getorbitalindex(args...) getsiteindex(args...)
@deprecate getorbitalcoord(args...) getsitecoord(args...)
@deprecate getorbitalindexcoord(args...) getsiteindexcoord(args...)
@deprecate getorbitalname(args...) getsitename(args...)
@deprecate findorbitalindex(args...; kwargs...) findsiteindex(args...; kwargs...)

"""
    UnitCell{O}

Represent a unitcell of a lattice, which contains sites at fixed locations
(does not yet implement multiple orbitals per site). It is recommended to use `makeunitcell`
rather than the constructor to make a UnitCell object.

# Parameters
- `O`: type of "site". Any type can be used, but we recommend using `String`
  or tuple of `String` and `Int` for compatibility with JSON.

# Fields
- `latticevectors::Array{Float64, 2}`: Lattice vectors
- `reducedreciprocallatticevectors::Array{Float64, 2}`: Reduced reciprocal lattice vectors
   (transpose of inverse of `latticevectors`)
- `reciprocallatticevectors ::Array{Float64, 2}`: Reciprocal lattice vectors.
  `2π * reducedreciprocallatticevectors`
- `sites::Vector{Tuple{T, FractCoord}}`: List of sites within unit cell
- `siteindices::Dict{T, Int}`: Indices of sites
"""
mutable struct UnitCell{S}
    latticevectors::Array{Float64, 2}
    sites::Vector{Tuple{S, FractCoord}}

    reducedreciprocallatticevectors::Array{Float64, 2}
    reciprocallatticevectors::Array{Float64, 2}
    siteindices::Dict{S, Int}

    function UnitCell{S}(
        latticevectors::AbstractArray{<:Real, 2},
        sites::AbstractVector{Tuple{S, FractCoord}},
        reducedreciprocallatticevectors::AbstractArray{<:Real, 2},
        reciprocallatticevectors::AbstractArray{<:Real, 2},
        siteindices::AbstractDict{S, Int}
    ) where {S}
        if (S <: Integer)
            throw(ArgumentError("Integer SiteType is disallowed to avoid confusion"))
        end
        new{S}(
            latticevectors,
            sites,
            reducedreciprocallatticevectors,
            reciprocallatticevectors,
            siteindices,
        )
    end
end

"""
    makeunitcell(latticevectors; SiteType=Any, tol=√ϵ)

Construct an n-dimensional lattice.

# Arguments
* `latticevectors`: Lattice vectors. Can be a nested array of lattice vectors,
  or a two-dimensional array whose columns are lattice vectors, or a real number.

# Optional Arguments
* `SiteType::DataType`
* `tol=√ϵ`: Epsilon
"""
function makeunitcell(
    latticevectors::AbstractArray{<:Real, 2};
    SiteType::DataType=Any,
    OrbitalType::DataType=Nothing,
    tol::Real=Base.rtoldefault(Float64)
)
    if OrbitalType != Nothing
        @warn "OrbitalType is deprecated. use SiteType instead"
        if SiteType == Any
            SiteType = OrbitalType
        else
            @error "Cannot specify OrbitalType and SiteType simultaneously"
        end
    end

    (ndim, ndim_) = size(latticevectors)
    if ndim != ndim_
        throw(ArgumentError("lattice vectors has dimension ($(ndim), $(ndim_))"))
    elseif abs(LinearAlgebra.det(latticevectors)) <= tol
        throw(ArgumentError("lattice vectors define zero volume $(latticevectors)"))
    elseif tol < 0
        throw(ArgumentError("tol must be non-negative"))
    end

    reduced_rlv = transpose(inv(latticevectors))
    sites = Tuple{SiteType, FractCoord}[]
    siteindices = Dict{SiteType, Int}()
    return UnitCell{SiteType}(latticevectors, sites, reduced_rlv, 2*π*reduced_rlv, siteindices)
end

function makeunitcell(
    latticevectors::AbstractVector{<:AbstractVector};
    SiteType::DataType=Any,
    tol::Real=Base.rtoldefault(Float64)
)
    lv = hcat(latticevectors...)
    return makeunitcell(lv; SiteType=SiteType, tol=tol)
end

function makeunitcell(
    latticeconstant::Real;
    SiteType::DataType=Any,
    OrbitalType::DataType=Nothing,
    tol::Real=Base.rtoldefault(Float64),
)
    return makeunitcell(
        reshape([latticeconstant], (1,1));
        SiteType=SiteType, OrbitalType=OrbitalType, tol=tol,
    )
end

make_unitcell = makeunitcell

"""
    dimension(uc)

Spatial dimension of the unit cell.
"""
dimension(uc::UnitCell) = size(uc.latticevectors, 1)

"""
    numsite(unitcell)

Number of sites of the unit cell.

# Arguments
* `uc ::UnitCell`
"""
numsite(uc::UnitCell) = length(uc.sites)

"""
    sitecount(unitcell)

Number of sites of the unit cell.

# Arguments
* `uc ::UnitCell`
"""
sitecount(uc::UnitCell) = length(uc.sites)

"""
    addsite!(unitcell, sitename, sitecoord)

Add an site to the unit cell.

# Arguments
* `uc ::UnitCell{T}`
* `sitename ::{T}`
* `sitecoord ::FractCoord`
"""
function addsite!(uc::UnitCell{O}, sitename::O, sitecoord::FractCoord) where {O}
    (ndim, ndim_) = size(uc.latticevectors)
    if dimension(sitecoord) != ndim
        throw(ArgumentError("sitecoord has wrong dimension"))
    elseif haskey(uc.siteindices, sitename)
        throw(ArgumentError( "duplicate site name"))
    end

    for site in uc.sites
        if isapprox(site[2].fraction, sitecoord.fraction)
            throw(ArgumentError("site $(site[1]) is already at $(sitecoord.fraction)"))
        end
    end

    push!(uc.sites, (sitename, sitecoord))
    index = length(uc.siteindices)+1
    uc.siteindices[sitename] = index
    return index
end

"""
    hassite(unitcell, name)

Test whether the unit cell contains the site of given name.

# Arguments
* `unitcell::UnitCell{O}`
* `name::O`
"""
function hassite(uc::UnitCell{O}, name::O) where {O}
    return haskey(uc.siteindices, name)
end

"""
    getsiteindex(unitcell, name)

Get index of the given site.

# Arguments
* `unitcell::UnitCell{O}`
* `name::O`
"""
function getsiteindex(uc::UnitCell{O}, name::O) where {O}
    return uc.siteindices[name]
end

"""
    getsite(unitcell, name)

Get the site (its site name and its fractional coordinates) with the given name.

# Arguments
* `unitcell::UnitCell{O}`
* `name::O`

# Return
* `(sitename, fractcoord)`
"""
function getsite(uc::UnitCell{O}, name::O) where {O}
    return uc.sites[uc.siteindices[name]]
end

"""
    getsitecoord(unitcell, name)

Get the fractional coordinates of the site with the given name.

# Arguments
* `uc ::UnitCell{O}`
* `name ::O`

# Return
* `fractcoord`
"""
function getsitecoord(uc::UnitCell{O}, name::O) where {O}
    return getsite(uc, name)[2]
end

"""
    getsiteindexcoord(unitcell, name)

# Arguments
* `unitcell::UnitCell{T}`
* `name::T`

# Return
* `(index, fractcoord)`
"""
function getsiteindexcoord(uc::UnitCell{O}, name::O) where {O}
    index = getsiteindex(uc, name)
    coord = getsitecoord(uc, index)
    return (index, coord)
end

"""
    getsite(unitcell, index)

# Arguments
* `unitcell::UnitCell`
* `index::Integer`

# Return
* `(sitename, fractcoord)`
"""
function getsite(uc::UnitCell, index::Integer)
    return uc.sites[index]
end

"""
    getsitename(unitcell, index)

# Arguments
* `uc::UnitCell`
* `index::Integer`

# Return
* `sitename`
"""
function getsitename(uc::UnitCell{O}, index::Integer)::O where {O}
    return uc.sites[index][1]
end

"""
    getsitecoord(unitcell, index)

# Arguments
* `unitcell::UnitCell`
* `index::Integer`

# Return
* `FractCoord`
"""
function getsitecoord(uc::UnitCell, index::Integer)::FractCoord
    return uc.sites[index][2]
end

"""
    fract2carte(unitcell, fractcoord)

# Arguments
* `unitcell::UnitCell`
* `fractcoord::FractCoord`
"""
function fract2carte(unitcell::UnitCell, fc::FractCoord)::CarteCoord
    if dimension(unitcell) != dimension(fc)
        throw(ArgumentError("unitcell and fractcoord must have the same dimension"))
    end
    mc = fc.whole + fc.fraction
    cc = unitcell.latticevectors * mc
    return CarteCoord(cc)
end

"""
    carte2fract(unitcell, cartecoord; tol=√ϵ)

# Arguments
* `unitcell::UnitCell`
* `cc::CarteCoord`
* `tol::Real=Base.rtoldefault(Float64)`
"""
function carte2fract(
    unitcell::UnitCell, cc::CarteCoord;
    tol::Real=Base.rtoldefault(Float64)
)
    if dimension(unitcell) != length(cc)
        throw(ArgumentError("unitcell and cartecoord must have the same dimension"))
    end
    fc = transpose(unitcell.reducedreciprocallatticevectors) * cc
    w = Int[fld(x, 1) for x in fc]
    f = Float64[mod(x, 1) for x in fc]
    return FractCoord(w, f; tol=tol)
end

"""
whichunitcell(unitcell, name, cartecoord; tol=√ϵ)

# Return
- `Vector{Int}`: the integer coordinate of the unitcell that the specified site/cartesian
  coordinate belongs to.
"""
function whichunitcell(
    uc::UnitCell{O},
    name::O,
    cc::CarteCoord;
    tol::Real=Base.rtoldefault(Float64)
) where {O}
    fc1 = getsitecoord(uc, name)
    fc2 = carte2fract(uc, cc)
    if !isapprox(fc1.fraction, fc2.fraction; atol=tol)
        throw(ArgumentError("$(fc1.fraction) != $(fc2.fraction) [tol=$(tol)]"))
    end
    R = fc2.whole - fc1.whole
    return R
end

"""
    whichunitcell(unitcell, name, fractcoord; tol=√ϵ)

# Return
- `Vector{Int}`: the integer coordinate of the unitcell that the specified site/fractional
  coordinate belongs to.
"""
function whichunitcell(
    uc::UnitCell{O},
    name::O,
    fc::FractCoord;
    tol::Real=Base.rtoldefault(Float64)
) where {O}
    fc1 = getsitecoord(uc, name)
    fc2 = fc
    if !isapprox(fc1.fraction, fc2.fraction; atol=tol)
        throw(ArgumentError("$(fc1.fraction) != $(fc2.fraction) [tol=$(tol)]"))
    end
    R = fc2.whole - fc1.whole
    return R
end

#=
function zero(uc::UnitCell; dtype::DataType=ComplexF64)
    norb = numsite(uc)
    return Base.zeros(dtype, (norb, norb))
end
=#

"""
    momentumgrid(unitcell, shape)

Generate an n-dimensional grid of momenta of given shape.

Returns an n-dimensional array of the following form:
```math
   k[i_1, i_2, \\ldots ] = G \\cdot ( \\frac{i_1}{n_1}, \\frac{i_2}{n_2}, \\ldots )
````
where G is the reciprocal lattice vector
"""
function momentumgrid(uc::UnitCell, shape::AbstractVector{<:Integer})
    if length(shape) != dimension(uc)
        throw(ArgumentError("dimension mismatch"))
    elseif !all((x) -> x>0, shape)
        throw(ArgumentError("shape should be positive"))
    end
    ranges = [range(0,stop=1,length=n+1)[1:end-1] for n in shape]
    cubicgrid = map((x) -> [x...], Base.product(ranges...))
    momentumgrid = map((x) -> uc.reciprocallatticevectors * x, cubicgrid)
    return momentumgrid
end

function Base.:(==)(lhs::UnitCell{O}, rhs::UnitCell{O}) where O
    return (
        lhs.latticevectors == rhs.latticevectors &&
        lhs.sites == rhs.sites
    )
end

"""
    findsiteindex(unitcell::UnitCell, fc::FractCoord; tol=√ϵ)

Find the index of the site, and the unitcell coordinate at the specified fractional coordinate.

# Arguments
- `unitcell::UnitCell`
- `fc::FractCoord`

# Returns
- `(site_index, unitcell_vector)`, or `(-1, [])` if not found.
"""
function findsiteindex(
    unitcell::UnitCell,
    fc::FractCoord;
    tol::Real=Base.rtoldefault(Float64)
)
    i = findfirst(
        x -> isapprox(x, fc.fraction; atol=tol),
        [orbfc.fraction for (orbname, orbfc) in unitcell.sites]
    )
    if i !== nothing
        (orbname, orbfc) = unitcell.sites[i]
        return (i, fc.whole - orbfc.whole)
    else
        return (-1, Int[])
    end
end


text:
##########################################################################
# Copyright 2017 Samuel Ridler.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##########################################################################

# load the resimulation data from the events file
# changes sim.resim.use to be false unless all checks are passed
function initResim!(sim::Simulation; doPrint::Bool = false)
	resim = sim.resim # shorthand
	@assert(resim.use)
	resim.use = false # until checks have passed

	eventsFilename = sim.outputFiles["events"].path
	doPrint && println("resimulating, based on events from file: ", eventsFilename)

	if !isfile(eventsFilename)
		doPrint && println("cannot resimulate, events file not found")
		return
	end

	# read events file
	(events, eventsChildren, eventFilter, fileEnded, inputFiles, fileChecksums) = readEventsFile(eventsFilename)

	if !fileEnded
		doPrint && println("cannot resimulate, events file closed before end")
		return
	end

	# check that checksum values of input files are same as in events file
	allMatch = true
	for i = 1:length(inputFiles)
		if inputFiles[i] == "statsControl" continue end # allow this file to change, should not affect sim events
		if fileChecksums[i] != sim.inputFiles[inputFiles[i]].checksum
			doPrint && println(" checksum mismatch for file: ", inputFiles[i])
			allMatch = false
		end
	end
	if !allMatch
		doPrint && println("cannot resimulate, input file checksums do not match those in events file")
		return
	end

	# all checks have passed, can resimulate
	resim.use = true
	resim.events = events
	resim.eventsChildren = eventsChildren
	resim.eventFilter = eventFilter
	resim.doDispatch = eventFilter[ambDispatched]
	resim.doMoveUp = eventFilter[ambMoveUpToStation]
	resim.prevEventIndex = 0
	resim.timeTolerance = 1e-5 / 2
end

function resimCheckCurrentEvent!(sim::Simulation, event::Event)
	resim = sim.resim # shorthand
	@assert(resim.use)

	i = resim.prevEventIndex # shorthand
	if i == length(resim.events) || resim.events[i+1].index > event.index
		return # have run out of events to resimulate, or do not have current event to resimulate with
	end
	resim.prevEventIndex += 1 # go to event after previous (should be current)

	resimEvent = resim.events[resim.prevEventIndex] # shorthand
	eventsMatch = true
	@assert(resimEvent.index == event.index)
	if !isapprox(event.time, resimEvent.time; rtol = eps(Float), atol = resim.timeTolerance)
		println("mismatching event time")
		eventsMatch = false
	elseif event.form != resimEvent.form
		println("mismatching event form")
		eventsMatch = false
	elseif event.ambIndex != resimEvent.ambIndex
		println("mismatching event ambulance index")
		eventsMatch = false
	elseif event.callIndex != resimEvent.callIndex
		println("mismatching event call index")
		eventsMatch = false
	elseif event.stationIndex != resimEvent.stationIndex
		println("mismatching event station index")
		eventsMatch = false
	end
	if !eventsMatch
		@show event
		@show resimEvent
		error("resimulation event does not match current event")
	end
end

# find ambulance to dispatch to a call
# assumes ambulance mobilisation delay is 0
function resimFindAmbToDispatch(sim::Simulation, call::Call)
	resim = sim.resim # shorthand
	@assert(resim.use && resim.doDispatch)
	eventIndex = sim.eventIndex # shorthand

	if resim.prevEventIndex > 0
		resimEvent = resim.events[resim.prevEventIndex]
		if resimEvent.index == eventIndex
			# check that previous event was to dispatch for this call
			@assert(resimEvent.form == considerDispatch)
			@assert(resimEvent.callIndex == call.index)
		end
	end

	# find child event for dispatch
	ambIndex = nullIndex
	if length(resim.eventsChildren) < eventIndex return ambIndex end
	eventChildren = resim.eventsChildren[eventIndex]
	dispatchEvents = filter(event -> event.form == ambDispatched, eventChildren)
	@assert(length(dispatchEvents) <= 1)
	if length(dispatchEvents) == 1
		ambIndex = dispatchEvents[1].ambIndex
	end

	return ambIndex
end

# move up function for re-simulation
# assumes ambulance move up delay is 0
function resimMoveUp(sim::Simulation)
	resim = sim.resim # shorthand
	@assert(resim.use && resim.doMoveUp)
	eventIndex = sim.eventIndex # shorthand

	if resim.prevEventIndex > 0
		resimEvent = resim.events[resim.prevEventIndex]
		if resimEvent.index == eventIndex
			# check that previous event was to consider move up
			@assert(resimEvent.form == considerMoveUp)
		end
	end

	# find all ambulances to move up
	movableAmbs = Vector{Ambulance}()
	ambStations = Vector{Station}()
	if length(resim.eventsChildren) < eventIndex return movableAmbs, ambStations end
	eventChildren = resim.eventsChildren[eventIndex]
	for event in reverse(eventChildren)
		@assert(event.form == ambMoveUpToStation)
		push!(movableAmbs, sim.ambulances[event.ambIndex])
		push!(ambStations, sim.stations[event.stationIndex])
	end

	return movableAmbs, ambStations
end


text:
module Tests
using AverageShiftedHistograms, Test, StatsBase, Statistics, LinearAlgebra

@info("Messy Output")
show(ash(randn(1000)))
show(ash(randn(100), randn(100)))
println("\n\n")
@info("Begin Tests")

@testset "Kernels" begin
    for kernel in [Kernels.biweight,
                   Kernels.cosine,
                   Kernels.epanechnikov,
                   Kernels.triangular,
                   Kernels.tricube,
                   Kernels.triweight,
                   Kernels.uniform,
                   Kernels.gaussian,
                   Kernels.logistic]
        for u in [0.0, 0.00001, 0.0001, 0.001, 0.01, 0.1, 1.0, 10.0]
            @test kernel(-u) ≈ kernel(u)
        end
    end
end

@testset "Ash" begin
    x = randn(10_000)
    o = ash(x, rng = -7:.01:7)
    @test extrema(o)[1] ≈ extrema(x)[1] atol=0.01
    @test extrema(o)[2] ≈ extrema(x)[2] atol=0.01

    for f in [mean, var, std, x -> quantile(x, .4), x -> quantile(x, .4:.1:.6)]
        @test f(o) ≈ f(x) atol=.1
    end
    AverageShiftedHistograms.histdensity(o)
    @test quantile(o, 0) ≈ o.rng[findnext(x -> x != 0, o.counts, 1)]

    # check that histogram is correct
    h = fit(Histogram, x, (-7:.01:7.01) .- .005; closed = :left)
    @test h.weights == o.counts

    ash!(o, x)
    @test nobs(o) == 20_000
    @test AverageShiftedHistograms.pdf(o, -5) == 0
    @test AverageShiftedHistograms.pdf(o, 5) == 0
    @test AverageShiftedHistograms.pdf(o, 0) > 0
    @test AverageShiftedHistograms.cdf(o, -5) == 0
    @test AverageShiftedHistograms.cdf(o, 0) ≈ .5 atol=.05
    @test AverageShiftedHistograms.cdf(o, 5) ≈ 1

    o = ash([.1, .1]; rng = -1:.1:1)
    @test nout(o) == 0

    # merging
    y = randn(100)
    y2 = randn(100)

    @test ash(y) == ash!(ash(y))

    a = ash(y; rng = -5:.1:5)
    ash!(a, y2)
    a2 = ash(vcat(y, y2); rng = -5:.1:5)
    @test a == a2
    a3 = merge(ash(y, rng=-5:.1:5), ash(y2, rng=-5:.1:5))
    @test a == a3

    # push!
    a = ash(y; rng = -5:.1:5)
    a2 = ash([y[1]]; rng = -5:.1:5)
    for i in 2:length(y)
        push!(a2, y[i])
    end
    ash!(a2)
    @test a == a2
end

@testset "Ash2" begin
    x = randn(1000)
    y = x + randn(1000)
    o = ash(x, y)

    @test ash(x, y) == ash!(ash(x, y))

    @test nobs(o) == 1000

    o = ash([1,2], [1,2])
    @test nout(o) == 0

    mean(o)
    var(o)
    std(o)
    xyz(o)
end

end #module


text:
"""
    DiscreteNonParametric(xs, ps)

A *Discrete nonparametric distribution* explicitly defines an arbitrary
probability mass function in terms of a list of real support values and their
corresponding probabilities

```julia
d = DiscreteNonParametric(xs, ps)

params(d)  # Get the parameters, i.e. (xs, ps)
support(d) # Get a sorted AbstractVector describing the support (xs) of the distribution
probs(d)   # Get a Vector of the probabilities (ps) associated with the support
```

External links

* [Probability mass function on Wikipedia](http://en.wikipedia.org/wiki/Probability_mass_function)
"""
struct DiscreteNonParametric{T<:Real,P<:Real,Ts<:AbstractVector{T},Ps<:AbstractVector{P}} <: DiscreteUnivariateDistribution
    support::Ts
    p::Ps

    function DiscreteNonParametric{T,P,Ts,Ps}(xs::Ts, ps::Ps; check_args::Bool=true) where {
            T<:Real,P<:Real,Ts<:AbstractVector{T},Ps<:AbstractVector{P}}
        check_args || return new{T,P,Ts,Ps}(xs, ps)
        @check_args(
            DiscreteNonParametric,
            (length(xs) == length(ps), "length of support and probability vector must be equal"),
            (ps, isprobvec(ps), "vector is not a probability vector"),
            (xs, allunique(xs), "support must contain only unique elements"),
        )
        sort_order = sortperm(xs)
        new{T,P,Ts,Ps}(xs[sort_order], ps[sort_order])
    end
end

DiscreteNonParametric(vs::AbstractVector{T}, ps::AbstractVector{P}; check_args::Bool=true) where {
        T<:Real,P<:Real} =
    DiscreteNonParametric{T,P,typeof(vs),typeof(ps)}(vs, ps; check_args=check_args)

Base.eltype(::Type{<:DiscreteNonParametric{T}}) where T = T

# Conversion
convert(::Type{DiscreteNonParametric{T,P,Ts,Ps}}, d::DiscreteNonParametric) where {T,P,Ts,Ps} =
    DiscreteNonParametric{T,P,Ts,Ps}(convert(Ts, support(d)), convert(Ps, probs(d)), check_args=false)

# Accessors
params(d::DiscreteNonParametric) = (d.support, d.p)

"""
    support(d::DiscreteNonParametric)

Get a sorted AbstractVector defining the support of `d`.
"""
support(d::DiscreteNonParametric) = d.support

"""
    probs(d::DiscreteNonParametric)

Get the vector of probabilities associated with the support of `d`.
"""
probs(d::DiscreteNonParametric)  = d.p

==(c1::D, c2::D) where D<:DiscreteNonParametric =
    (support(c1) == support(c2) || all(support(c1) .== support(c2))) &&
    (probs(c1) == probs(c2) || all(probs(c1) .== probs(c2)))

Base.isapprox(c1::D, c2::D) where D<:DiscreteNonParametric =
    (support(c1) ≈ support(c2) || all(support(c1) .≈ support(c2))) &&
    (probs(c1) ≈ probs(c2) || all(probs(c1) .≈ probs(c2)))

# Sampling

function rand(rng::AbstractRNG, d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    n = length(p)
    draw = rand(rng, float(eltype(p)))
    cp = p[1]
    i = 1
    while cp <= draw && i < n
        @inbounds cp += p[i +=1]
    end
    return x[i]
end

sampler(d::DiscreteNonParametric) =
    DiscreteNonParametricSampler(support(d), probs(d))

# Override the method in testutils.jl since it assumes
# an evenly-spaced integer support
get_evalsamples(d::DiscreteNonParametric, ::Float64) = support(d)

# Evaluation

pdf(d::DiscreteNonParametric) = copy(probs(d))

function pdf(d::DiscreteNonParametric, x::Real)
    s = support(d)
    idx = searchsortedfirst(s, x)
    ps = probs(d)
    if idx <= length(ps) && s[idx] == x
        return ps[idx]
    else
        return zero(eltype(ps))
    end
end
logpdf(d::DiscreteNonParametric, x::Real) = log(pdf(d, x))

function cdf(d::DiscreteNonParametric, x::Real)
    ps = probs(d)
    P = float(eltype(ps))

    # trivial cases
    x < minimum(d) && return zero(P)
    x >= maximum(d) && return one(P)
    isnan(x) && return P(NaN)

    n = length(ps)
    stop_idx = searchsortedlast(support(d), x)
    s = zero(P)
    if stop_idx < div(n, 2)
        @inbounds for i in 1:stop_idx
            s += ps[i]
        end
    else
        @inbounds for i in (stop_idx + 1):n
            s += ps[i]
        end
        s = 1 - s
    end

    return s
end

function ccdf(d::DiscreteNonParametric, x::Real)
    ps = probs(d)
    P = float(eltype(ps))

    # trivial cases
    x < minimum(d) && return one(P)
    x >= maximum(d) && return zero(P)
    isnan(x) && return P(NaN)

    n = length(ps)
    stop_idx = searchsortedlast(support(d), x)
    s = zero(P)
    if stop_idx < div(n, 2)
        @inbounds for i in 1:stop_idx
            s += ps[i]
        end
        s = 1 - s
    else
        @inbounds for i in (stop_idx + 1):n
            s += ps[i]
        end
    end

    return s
end

function quantile(d::DiscreteNonParametric, q::Real)
    0 <= q <= 1 || throw(DomainError())
    x = support(d)
    p = probs(d)
    k = length(x)
    i = 1
    cp = p[1]
    while cp < q && i < k #Note: is i < k necessary?
        i += 1
        @inbounds cp += p[i]
    end
    x[i]
end

minimum(d::DiscreteNonParametric) = first(support(d))
maximum(d::DiscreteNonParametric) = last(support(d))
insupport(d::DiscreteNonParametric, x::Real) =
    length(searchsorted(support(d), x)) > 0

mean(d::DiscreteNonParametric) = dot(probs(d), support(d))

function var(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return var(x, Weights(p, one(eltype(p))); corrected=false)
end

function skewness(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return skewness(x, Weights(p, one(eltype(p))))
end

function kurtosis(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return kurtosis(x, Weights(p, one(eltype(p))))
end

entropy(d::DiscreteNonParametric) = entropy(probs(d))
entropy(d::DiscreteNonParametric, b::Real) = entropy(probs(d), b)

function mode(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return mode(x, Weights(p, one(eltype(p))))
end
function modes(d::DiscreteNonParametric)
    x = support(d)
    p = probs(d)
    return modes(x, Weights(p, one(eltype(p))))
end

function mgf(d::DiscreteNonParametric, t::Real)
    x, p = params(d)
    s = zero(Float64)
    for i in 1:length(x)
        s += p[i] * exp(t*x[i])
    end
    s
end

function cf(d::DiscreteNonParametric, t::Real)
    x, p = params(d)
    s = zero(Complex{Float64})
    for i in 1:length(x)
       s += p[i] * cis(t*x[i])
    end
    s
end

# Sufficient statistics

struct DiscreteNonParametricStats{T<:Real,W<:Real,Ts<:AbstractVector{T},
                                  Ws<:AbstractVector{W}} <: SufficientStats
    support::Ts
    freq::Ws
end

function suffstats(::Type{<:DiscreteNonParametric}, x::AbstractArray{T}) where {T<:Real}

    N = length(x)
    N == 0 && return DiscreteNonParametricStats(T[], Float64[])

    n = 1
    vs = Vector{T}(undef,N)
    ps = zeros(Float64, N)
    x = sort(vec(x))

    vs[1] = x[1]
    ps[1] += 1.

    xprev = x[1]
    @inbounds for i = 2:N
        xi = x[i]
        if xi != xprev
            n += 1
            vs[n] = xi
        end
        ps[n] += 1.
        xprev = xi
    end

    resize!(vs, n)
    resize!(ps, n)
    DiscreteNonParametricStats(vs, ps)

end

function suffstats(::Type{<:DiscreteNonParametric}, x::AbstractArray{T},
                   w::AbstractArray{W}; check_args::Bool=true) where {T<:Real,W<:Real}

    @check_args DiscreteNonParametric (length(x) == length(w))

    N = length(x)
    N == 0 && return DiscreteNonParametricStats(T[], W[])

    n = 1
    vs = Vector{T}(undef, N)
    ps = zeros(W, N)

    xorder = sortperm(vec(x))
    x = vec(x)[xorder]
    w = vec(w)[xorder]

    vs[1] = x[1]
    ps[1] += w[1]

    xprev = x[1]
    @inbounds for i = 2:N
        xi = x[i]
        wi = w[i]
        if xi != xprev
            n += 1
            vs[n] = xi
        end
        ps[n] += wi
        xprev = xi
    end

    resize!(vs, n)
    resize!(ps, n)
    DiscreteNonParametricStats(vs, ps)

end

# # Model fitting

fit_mle(::Type{<:DiscreteNonParametric},
        ss::DiscreteNonParametricStats{T,W,Ts,Ws}) where {T,W,Ts,Ws} =
    DiscreteNonParametric{T,W,Ts,Ws}(ss.support, normalize!(copy(ss.freq), 1), check_args=false)


text:
export plot_anymod_result

pjs_rgb(r::Int,g::Int,b::Int) = "rgb($r, $g, $b))"
rand_rgb() = pjs_rgb(rand(1:255), rand(1:255), rand(1:255))

function get_x_axis(row::AbstractArray)
    return map(x-> x.colname |> Symbol, row)
end

get_x_axis(row) = row.colname |> Symbol

xval(df, x::AbstractArray) = [df[!,x[1]], df[!,x[2]]]
xval(df, x::Symbol) = df[!,x]

function grouped_plot(df::AbstractDataFrame, x, col; colors::Dict=Dict())
    g = col.colname |> Symbol
    g_vals = unique(df[:,g])
    traces = map(g_vals) do g_key
        df_g = filter(y-> y[g] == g_key, df)
        t = GenericTrace(
            xval(df_g, x),
            df_g[!, :value],
            kind = "bar",
            name = g_key,
            color = get(colors, g_key, rand_rgb())
        )
        return t
    end

    return traces
end

function grouped_plot(df::AbstractDataFrame, x, col::Nothing;
    colors::Dict=Dict()
)
    trace = GenericTrace(
        xval(df, x),
        df[!, :value],
        kind = "bar"
    )

    return trace
end

function plot_anymod_result(ar::AnymodResult, m::Mask)
    df = stackedresult(ar, m)
    x = get_x_axis(m.row)
    traces = grouped_plot(df, x, m.col)
    layout = Layout(;barmode="relative")
    return Plot(traces, layout)
end


text:
using DiseaseOutbreak
using Plots

#####################################
# more infectious
#####################################
n=10000
status,d,pos = makeCitizens(n);
#TODO: confirm with Saunak: infectCitizens needs second argument: nInfect
infectCitizens!(status, 10)
followCitizens!(status,d,365,0.05,14,0.0025)

uninfected(status)

newInfections(status)

cumulativeInfections(status)

gif(infectionSpread(status,pos),fps=15)

#####################################
# less infectious
#####################################
n=10000
status1,d1,pos1 = makeCitizens(n);
infectCitizens!(status1, 10)
followCitizens!(status1,d1,365,0.05,14,0.002)
uninfected(status1)

newInfections(status1)

cumulativeInfections(status1)

gif(infectionSpread(status1,pos1),fps=15)


text:
using JLD
using BenchmarkTools
using Dates
using Base.Sys
using CellListsBenchmarks
using Plots

gettime(x) = getfield(x, :time)
gettime_nogc(x) = getfield(x, :time) - getfield(x, :gctime)

filepath_serial = "output/near_neighbors_serial_2021-04-05T11:47:02+0300/n20000-d2-r0.01-nthreads1.jld"
filepath_parallel = "output/near_neighbors_parallel_2021-04-05T11:47:16+0300/n20000-d2-r0.01-nthreads4.jld"

ls = load(filepath_serial)["benchmarks"]
lp = load(filepath_parallel)["benchmarks"]

d = getfield(ls[1], :d)
r = getfield(ls[1], :r)
n = getfield(ls[1], :n)
ts = getfield(ls[1], :trials)
tp = getfield(lp[1], :trials)

ms = gettime_nogc.(median.(ts))
mp = gettime_nogc.(median.(tp))

p1 = plot(legend=false)
bar!(p1, ms, alpha=0.5)
bar!(p1, mp, alpha=0.5)

rt = ms ./ mp
p2 = scatter(rt, legend=false)
plot!(p2, 1:length(rt), fill(mean(rt), length(rt)))

directory = joinpath("figures", "near_neighbors")
if !isdir(directory)
    mkpath(directory)
end

savefig(p1, joinpath(directory, "d$d-r$r-n$n.svg"))
savefig(p2, joinpath(directory, "d$d-r$r-n$n-ratio.svg"))

max_stars_repo_name: jaantollander/cell-lists-benchmarks

max_issues_repo_name: jaantollander/cell-lists-benchmarks

max_forks_repo_name: jaantollander/cell-lists-benchmarks


text:
@testset "Example α_offset" begin

    α0=LinearFilter(Float64[0,1,0],0)
    α1=LinearFilter(Float64[0,1,0],1)
    β=collect(Float64,1:6)
    γ1=directConv(α0,1,β,ZeroPaddingBE,ZeroPaddingBE)
    γ2=directConv(α1,1,β,ZeroPaddingBE,ZeroPaddingBE)

    @test γ1 ≈ [2.0, 3.0, 4.0, 5.0, 6.0, 0.0]
    @test γ2 ≈ [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]

end;

@testset "Adjoint operator" begin

    α=LinearFilter(rand(4),2);
    β=rand(10);

    vβ=rand(length(β))
    d1=dot(directConv(α,-3,vβ,ZeroPaddingBE,ZeroPaddingBE),β)
    d2=dot(directConv(α,+3,β,ZeroPaddingBE,ZeroPaddingBE),vβ)

    @test isapprox(d1,d2)

    d1=dot(directConv(α,-3,vβ,PeriodicBE,PeriodicBE),β)
    d2=dot(directConv(α,+3,β,PeriodicBE,PeriodicBE),vβ)

    @test isapprox(d1,d2)

end;

@testset "Convolution commutativity" begin

    α=rand(4);
    αf=LinearFilter(α,0)
    β=rand(10);
    βf=LinearFilter(β,0)
    v1=zeros(20)
    v2=zeros(20)
    directConv!(αf,-1,
                β,v1,UnitRange(1,20),ZeroPaddingBE,ZeroPaddingBE)
    directConv!(βf,-1,
                α,v2,UnitRange(1,20),ZeroPaddingBE,ZeroPaddingBE)

    @test isapprox(v1,v2)

end;

@testset "Interval split" begin

    α=LinearFilter(rand(4),3)
    β=rand(10);

    γ=directConv(α,2,β,MirrorBE,PeriodicBE) # global computation
    Γ=zeros(length(γ))
    Ω1=UnitRange(1:3)
    Ω2=UnitRange(4:length(γ))
    directConv!(α,2,β,Γ,Ω1,MirrorBE,PeriodicBE) # compute on Ω1
    directConv!(α,2,β,Γ,Ω2,MirrorBE,PeriodicBE) # compute on Ω2

    @test isapprox(γ,Γ)

end;

@testset "2D convolution" begin
    β=rand(5,8)
    β_save=deepcopy(β)

    α_I=LinearFilter(Float64[0,2,0])
    α_J=LinearFilter(Float64[0,0,2,0,0])

    directConv2D!(α_I,1,α_J,-1,β)

    @test isapprox(β,4*β_save)
end;

@testset "2D crossCorrelation" begin
    β=rand(5,8)

    α_I=LinearFilter(Float64[0,2,0])
    α_J=LinearFilter(Float64[0,0,2,0,0])

    γ=directCrossCorrelation2D(α_I,α_J,β)

    @test isapprox(γ,4*β)
end;

max_stars_repo_name: vincent-picaud/DirectConvolution

max_issues_repo_name: vincent-picaud/DirectConvolution

max_forks_repo_name: vincent-picaud/DirectConvolution


text:
## Copyright (c) 2013 Miguel Bazdresch
##
## This file is distributed under the 2-clause BSD License.

# Auxiliary, non-exported functions are declared here.

# close gnuplot pipes
function gnuplot_exit(x...)
    for p in [gstdin, gstdout, gstderr]
        close(p)
    end
end

# Return index to figure with handle `c`. If no such figure exists, returns 0.
function findfigure(c)
    global gnuplot_state
    i = 0
    for j = 1:length(gnuplot_state.figs)
        if gnuplot_state.figs[j].handle == c
            i = j
            break
        end
    end
    return i
end

# remove a figure's data without closing it
function clearfigure(h::Int)
    global gnuplot_state

    f = findfigure(h)
    if f != 0
        gnuplot_state.figs[f] = Figure(h)
    end
end

# Return the next available handle (smallest non-used positive integer)
function nexthandle()
    isempty(gnuplot_state.figs) && return 1
    handles = [f.handle for f in gnuplot_state.figs]
    mh = maximum(handles)
    for i = 1:mh+1
        !in(i,handles) && return i
    end
end

# Push configuration, axes or curves to a figure. The handle is assumed valid.
function push_figure!(handle,args...)
    index = findfigure(handle)
    for c in args
        if isa(c,AxesConf)
            gnuplot_state.figs[index].conf = c
        elseif isa(c, Curve)
            if gnuplot_state.figs[index].isempty
                gnuplot_state.figs[index].curves = [c]
            else
                push!(gnuplot_state.figs[index].curves,c)
            end
            gnuplot_state.figs[index].isempty = false
        elseif isa(c, AbstractString)
            gnuplot_state.figs[index].gpcom = c
        end
    end
end

# convert marker string description to gnuplot's expected number
function pointtype(x::AbstractString)
    if x == "+"
        return 1
    elseif x == "x"
        return 2
    elseif x == "*"
        return 3
    elseif x == "esquare"
        return 4
    elseif x == "fsquare"
        return 5
    elseif x == "ecircle"
        return 6
    elseif x == "fcircle"
        return 7
    elseif x == "etrianup"
        return 8
    elseif x == "ftrianup"
        return 9
    elseif x == "etriandn"
        return 10
    elseif x == "ftriandn"
        return 11
    elseif x == "edmd"
        return 12
    elseif x == "fdmd"
        return 13
    end
    return 1
end

# create a Z-coordinate matrix from x, y coordinates and a function
function meshgrid(x,y,f)
    Z = zeros(length(x),length(y))
    for k = 1:length(x)
        Z[k,:] = [ f(i,j) for i=x[k], j=y ]
    end
    return Z
end

# create x,y coordinates for a histogram, from a sample vector, using a number
# of bins
function hist(s,bins)
    # When adding an element s to a bin, we use an iequality m < s <= M.
    # In order to account for elements s==m, we need to special-case
    # the computation for the first bin
    ms = minimum(s)
    Ms = maximum(s)
    bins = max(bins, 1)
    if Ms == ms
        # compute a "natural" scale
        g = (10.0^floor(log10(abs(ms)+eps()))) / 2
        ms, Ms = ms - g, ms + g
    end
    delta = (Ms-ms)/bins
    x = ms:delta:Ms
    y = zeros(bins)
    # this is special-cased because we want to include the minimum in the
    # first bin
    y[1] = sum(ms .<= s .<= x[2])
    for i in 2:length(x)-2
        y[i] = sum(x[i] .< s .<= x[i+1])
    end
    # this is special-cased because there is no guarantee that x[end] == Ms
    # (because of how ranges work)
    if length(y) > 1 y[end] = sum(x[end-1] .< s .<= Ms) end

    if bins != 1
        # We want the left bin to start at ms and the right bin to end at Ms
        x = (ms+delta/2):delta:Ms
    else
        # add two empty bins on the sides to provide a scale to gnuplot
        x = (ms-delta/2):delta:(ms+delta/2)
        y = [0.0, y[1], 0.0]
    end
    return x,y
end

function Base.show(io::IO, ::MIME"text/plain", x::Figure)
    if !isjupyter
        llplot()
        if gaston_config.terminal == "dumb" || gaston_config.terminal == "sixelgd"
            print(x.svg)
        end
        return nothing
    end
end

function Base.show(io::IO, ::MIME"image/svg+xml", x::Figure)
    llplot()
    write(io,x.svg)
end

# return configuration string for a single plot
function linestr_single(conf::CurveConf)
    s = ""
    if conf.legend != ""
        s = string(s, " title '", conf.legend, "' ")
    else
        s = string(s, "notitle ")
    end
    s = string(s, " with ", conf.plotstyle, " ")
    if conf.color != ""
        s = string(s, "linecolor rgb '", conf.color, "' ")
    end
    s = string(s, "lw ", string(conf.linewidth), " ")
    if conf.linestyle != ""
        s = string(s, "dt '", string(conf.linestyle), "' ")
    end
    # some plotstyles don't allow point specifiers
    cp = conf.plotstyle
    if cp != "lines" && cp != "impulses" && cp != "pm3d" && cp != "image" &&
        cp != "rgbimage" && cp != "boxes" && cp != "dots" && cp != "steps" &&
        cp != "fsteps" && cp != "fillsteps" && cp != "financebars"
        if conf.marker != ""
            s = string(s, "pt ", string(pointtype(conf.marker)), " ")
        end
        s = string(s, "ps ", string(conf.pointsize), " ")
    end
    return s
end

# build a string with plot commands according to configuration
function linestr(curves::Vector{Curve}, cmd::AbstractString, file::AbstractString)
    # We have to insert "," between plot commands. One easy way to do this
    # is create the first plot command, then the rest
    # We also need to keep track of the current index (starts at zero)
    index = 0
    s = string(cmd, " '", file, "' ", " i 0 ", linestr_single(curves[1].conf))
    if length(curves) > 1
        for i in curves[2:end]
            index += 1
            s = string(s, ", '", file, "' ", " i ", string(index)," ",
                       linestr_single(i.conf))
        end
    end
    return s
end

# Build a "set term" string appropriate for the terminal type
function termstring()
    global gnuplot_state
    global gaston_config

    gc = gaston_config
    term = gaston_config.terminal
    termname = term
    term == "null" && (termname = "dumb")
    term == "ijulia" && (termname = "svg")
    term == "pdf" && (termname = "pdfcairo")
    term == "eps" && (termname = "epscairo")

    ts = "set term $termname "

    if term ∈ term_window
        ts = ts*string(gnuplot_state.current)*" "
    end

    if term ∈ term_sup_font
        ts = ts*" font \"$(gc.print_fontface),$(gc.print_fontsize)\" "
    end

    if term ∈ term_sup_fontscale
        ts = ts*" fontscale $(gc.print_fontscale) "
    end

    if term ∈ term_sup_lw
        ts = ts*" linewidth $(gc.print_linewidth) "
    end

    if term ∈ term_sup_size
        if gc.print_size == ""  # use appropriate default size
            term ∈ term_size_in && (size = gc.print_size_in)
            term ∈ term_size_pix && (size = gc.print_size_pix)
        else
            size = gc.print_size  # use user-provided size
        end
        ts = ts*" size $size "
    end

    if term ∈ term_file
        # verify that user has set an output file
        gc.outputfile == "" && error("Plotting to file, but no file name given. Use `set(outputfile=\"filename\")` to configure the output file name.")
        ts = ts*"\nset output \"$(gc.outputfile)\" "
    end

    return ts
end

# send gnuplot the current figure's configuration
function gnuplot_send_fig_config(config)
    # fill style
    if config.fill != ""
        gnuplot_send(string("set style fill ",config.fill))
    end
    # grid
    if config.grid != ""
        if config.grid == "on"
            gnuplot_send(string("set grid"))
        else
            gnuplot_send(string("unset grid"))
        end
    end
    # legend box
    if config.box != ""
        gnuplot_send(string("set key ",config.box))
    end
    # plot title
    if config.title != ""
        gnuplot_send(string("set title '",config.title,"' "))
    end
    # xlabel
    if config.xlabel != ""
        gnuplot_send(string("set xlabel '",config.xlabel,"' "))
    end
    # ylabel
    if config.ylabel != ""
        gnuplot_send(string("set ylabel '",config.ylabel,"' "))
    end
    # zlabel
    if config.zlabel != ""
        gnuplot_send(string("set zlabel '",config.zlabel,"' "))
    end
    # axis log scale
    if config.axis != "" || config.axis != "normal"
        if config.axis == "semilogx"
            gnuplot_send("set logscale x")
        end
        if config.axis == "semilogy"
            gnuplot_send("set logscale y")
        end
        if config.axis == "loglog"
            gnuplot_send("set logscale xy")
        end
    end
    # ranges
    gnuplot_send("set autoscale")
    if config.xrange != ""
        gnuplot_send(string("set xrange ",config.xrange))
    end
    if config.yrange != ""
        gnuplot_send(string("set yrange ",config.yrange))
    end
    if config.zrange != ""
        gnuplot_send(string("set zrange ",config.zrange))
    end

    if config.palette != ""
        gnuplot_send(string("set palette ",config.palette))
    end
end

# write commands to gnuplot's pipe
function gnuplot_send(s::AbstractString)
    w = write(gstdin, string(s,"\n"))
    # check that data was accepted by the pipe
    if !(w > 0)
        println("Something went wrong writing to gnuplot STDIN.")
        return
    end
    flush(gstdin)
end


text:
using DFTK
using Test

include("testcases.jl")

function test_kernel(spin_polarization, termtype; test_compute=true)
    kgrid = [2, 2, 2]
    testcase = silicon
    ε   = 1e-8
    tol = 1e-5

    xcsym = (termtype isa Xc) ? join(string.(termtype.functionals), " ") : ""
    @testset "Kernel $(typeof(termtype)) $xcsym ($spin_polarization)" begin
        spec = ElementPsp(testcase.atnum, psp=load_psp(testcase.psp))
        magnetic_moments = []
        n_spin = 1
        if spin_polarization == :collinear
            magnetic_moments = [spec => 2rand(2)]
            n_spin = 2
        end

        model = Model(testcase.lattice; atoms=[spec => testcase.positions],
                      terms=[termtype], magnetic_moments=magnetic_moments,
                      spin_polarization=spin_polarization)
        @test model.n_spin_components == n_spin
        basis = PlaneWaveBasis(model; Ecut=2, kgrid=kgrid)
        term  = only(basis.terms)

        ρ0 = guess_density(basis, magnetic_moments)
        δρ = randn(size(ρ0))
        ρ_minus     = ρ0 - ε * δρ
        ρ_plus      = ρ0 + ε * δρ
        ops_minus = DFTK.ene_ops(term, basis, nothing, nothing; ρ=ρ_minus).ops
        ops_plus  = DFTK.ene_ops(term, basis, nothing, nothing; ρ=ρ_plus).ops
        δV = zero(ρ0)

        for iσ in 1:model.n_spin_components
            # Index of the first spin-up or spin-down k-point
            ifirst = first(DFTK.krange_spin(basis, iσ))
            δV[:, :, :, iσ] = (ops_plus[ifirst].potential - ops_minus[ifirst].potential) / (2ε)
        end

        δV_apply = DFTK.apply_kernel(term, basis, δρ; ρ=ρ0)
        @test norm(δV - δV_apply) < tol
        if test_compute
            kernel = DFTK.compute_kernel(term, basis; ρ=ρ0)
            δV_matrix = reshape(kernel * vec(δρ), size(δρ))
            @test norm(δV - δV_matrix) < tol
        end
    end
end

function test_kernel_collinear_vs_noncollinear(termtype)
    Ecut=2
    kgrid = [2, 2, 2]
    testcase = silicon

    xcsym = (termtype isa Xc) ? join(string.(termtype.functionals), " ") : ""
    @testset "Kernel $(typeof(termtype)) $xcsym (coll == noncoll)" begin
        spec = ElementPsp(testcase.atnum, psp=load_psp(testcase.psp))
        model = Model(testcase.lattice; atoms=[spec => testcase.positions],
                      terms=[termtype])
        basis = PlaneWaveBasis(model; Ecut, kgrid)
        term  = only(basis.terms)

        model_col = Model(testcase.lattice; atoms=[spec => testcase.positions],
                          terms=[termtype], spin_polarization=:collinear)
        basis_col = PlaneWaveBasis(model_col; Ecut, kgrid)
        term_col  = only(basis_col.terms)

        ρ0 = guess_density(basis)
        δρ = randn(size(ρ0))
        δV = DFTK.apply_kernel(term, basis, δρ; ρ=ρ0)

        ρ0_col = cat(0.5ρ0, 0.5ρ0, dims=4)
        δρ_col = cat(0.5δρ, 0.5δρ, dims=4)
        δV_pol = DFTK.apply_kernel(term_col, basis_col, δρ_col; ρ=ρ0_col)

        @test norm(δV_pol[:, :, :, 1] - δV_pol[:, :, :, 2]) < 1e-12
        @test norm(δV - δV_pol[:, :, :, 1:1]) < 1e-11
    end
end

test_kernel(:none, PowerNonlinearity(1.2, 2.0))
test_kernel(:none, Hartree())
test_kernel(:none, Xc(:lda_xc_teter93))
test_kernel(:none, Xc(:gga_c_pbe), test_compute=false)
test_kernel(:none, Xc(:gga_x_pbe), test_compute=false)

test_kernel_collinear_vs_noncollinear(Hartree())
test_kernel_collinear_vs_noncollinear(Xc(:lda_xc_teter93))
test_kernel_collinear_vs_noncollinear(Xc(:gga_c_pbe))
test_kernel_collinear_vs_noncollinear(Xc(:gga_x_pbe))

test_kernel(:collinear, Hartree())
test_kernel(:collinear, PowerNonlinearity(1.2, 2.5))
test_kernel(:collinear, Xc(:lda_xc_teter93))
test_kernel(:collinear, Xc(:gga_c_pbe), test_compute=false)
test_kernel(:collinear, Xc(:gga_x_pbe), test_compute=false)
test_kernel(:collinear, Xc(:gga_x_pbe, :gga_c_pbe), test_compute=false)


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct ContainerServiceVMDiagnostics <: SwaggerModel
    enabled::Any # spec type: Union{ Nothing, Bool } # spec name: enabled
    storageUri::Any # spec type: Union{ Nothing, String } # spec name: storageUri

    function ContainerServiceVMDiagnostics(;enabled=nothing, storageUri=nothing)
        o = new()
        validate_property(ContainerServiceVMDiagnostics, Symbol("enabled"), enabled)
        setfield!(o, Symbol("enabled"), enabled)
        validate_property(ContainerServiceVMDiagnostics, Symbol("storageUri"), storageUri)
        setfield!(o, Symbol("storageUri"), storageUri)
        o
    end
end # type ContainerServiceVMDiagnostics

const _property_map_ContainerServiceVMDiagnostics = Dict{Symbol,Symbol}(Symbol("enabled")=>Symbol("enabled"), Symbol("storageUri")=>Symbol("storageUri"))
const _property_types_ContainerServiceVMDiagnostics = Dict{Symbol,String}(Symbol("enabled")=>"Bool", Symbol("storageUri")=>"String")
Base.propertynames(::Type{ ContainerServiceVMDiagnostics }) = collect(keys(_property_map_ContainerServiceVMDiagnostics))
Swagger.property_type(::Type{ ContainerServiceVMDiagnostics }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_ContainerServiceVMDiagnostics[name]))}
Swagger.field_name(::Type{ ContainerServiceVMDiagnostics }, property_name::Symbol) =  _property_map_ContainerServiceVMDiagnostics[property_name]

function check_required(o::ContainerServiceVMDiagnostics)
    (getproperty(o, Symbol("enabled")) === nothing) && (return false)
    true
end

function validate_property(::Type{ ContainerServiceVMDiagnostics }, name::Symbol, val)
end


text: # anova.jl


text:
# Note that this script can accept some limited command-line arguments, run
# `julia build_tarballs.jl --help` to see a usage message.
using BinaryBuilder, Pkg

name = "GRAMPC"
version = v"2.2"

# The only Julia function used in this library is jl_error, so a version restriction is
# only needed because we need to have one to figure out the allowed platforms to build on.
# In reality, this library should work with other versions without the need to recompile it.
julia_version = v"1.6.0"

# Collection of sources required to complete build
sources = [
    GitSource( "https://github.com/JuDO-dev/grampc_julia_wrapper.git", "93f13feb53f4d27d92fcb29be61d34de0ade8cdf" )
]

# Bash recipe for building across all platforms
script = raw"""
    cd $WORKSPACE/srcdir/grampc_julia_wrapper

    mkdir -p build
    cd build

    cmake -DCMAKE_INSTALL_PREFIX=${prefix} \
          -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TARGET_TOOLCHAIN} \
          -DCMAKE_C_FLAGS="-I${prefix}/include" \
          -DCMAKE_BUILD_TYPE=Release \
          ../

    make -j${nproc}
    make install

    # Copy the library's license file
    install_license $WORKSPACE/srcdir/grampc_julia_wrapper/src_GRAMPC/LICENSE.txt
"""

# These are the platforms we will build for by default, unless further
# platforms are passed in on the command line
# Since we require the Julia library for linking, we can only build on the
# platforms it exists on.
include( "../../L/libjulia/common.jl" )
platforms = libjulia_platforms( julia_version )

# The products that we will ensure are always built
products = [
    LibraryProduct("libgrampc_julia", :libgrampc_julia)
]

# Dependencies that must be installed before this package can be built
dependencies = [
    BuildDependency( PackageSpec( name="libjulia_jll", version=julia_version ) ),
]

# Build the tarballs, and possibly a `build.jl` as well.
build_tarballs(ARGS, name, version, sources, script, platforms, products, dependencies)

max_stars_repo_name: sharanry/Yggdrasil

max_issues_repo_name: sharanry/Yggdrasil

max_forks_repo_name: sharanry/Yggdrasil


text:
#####
##### CPU horizontally periodic pressure solver
#####

function HorizontallyPeriodicPressureSolver(::CPU, grid, pressure_bcs, planner_flag=FFTW.PATIENT)
    kx², ky², kz² = generate_discrete_eigenvalues(grid, pressure_bcs)
    wavenumbers = (kx² = kx², ky² = ky², kz² = kz²)

    # Storage for RHS and Fourier coefficients is hard-coded to be Float64
    # because of precision issues with Float32.
    # See https://github.com/climate-machine/Oceananigans.jl/issues/55
    storage = zeros(Complex{Float64}, grid.Nx, grid.Ny, grid.Nz)

    @debug "Planning transforms for PressureSolver{HorizontallyPeriodic, CPU}..."
    x_bc, y_bc, z_bc = pressure_bcs.x.left, pressure_bcs.y.left, pressure_bcs.z.left
    FFTxy!  = plan_forward_transform(storage, x_bc, [1, 2], planner_flag)
    DCTz!   = plan_forward_transform(storage, z_bc, 3, planner_flag)
    IFFTxy! = plan_backward_transform(storage, x_bc, [1, 2], planner_flag)
    IDCTz!  = plan_backward_transform(storage, z_bc, 3, planner_flag)
    @debug "Planning transforms for PressureSolver{HorizontallyPeriodic, CPU} done!"

    transforms = ( FFTxy! =  FFTxy!,  DCTz! =  DCTz!,
                  IFFTxy! = IFFTxy!, IDCTz! = IDCTz!)

    return PressureSolver(HorizontallyPeriodic(), CPU(), wavenumbers, storage, transforms, nothing)
end

"""
    solve_poisson_equation!(solver::PressureSolver{HorizontallyPeriodic, CPU}, grid)

Solve Poisson equation on a uniform staggered grid (Arakawa C-grid) with
appropriate boundary conditions as specified by `solver.bcs` using planned FFTs
and DCTs. The right-hand-side RHS is stored in solver.storage which the solver
mutates to produce the solution, so it will also be stored in solver.storage.

We should describe the algorithm in detail in the documentation.
"""
function solve_poisson_equation!(solver::PressureSolver{HorizontallyPeriodic, CPU}, grid)
    Nx, Ny, Nz, _ = unpack_grid(grid)
    kx², ky², kz² = solver.wavenumbers.kx², solver.wavenumbers.ky², solver.wavenumbers.kz²

    # We can use the same storage for the RHS and the solution ϕ.
    RHS, ϕ = solver.storage, solver.storage

    solver.transforms.DCTz!  * RHS  # Calculate DCTᶻ(f) in place.
    solver.transforms.FFTxy! * RHS  # Calculate FFTˣʸ(f) in place.

    # Solve the discrete Poisson equation in spectral space. We are computing
    # the Fourier coefficients of the solution from the Fourier coefficients
    # of the RHS.
    @. ϕ = -RHS / (kx² + ky² + kz²)

    # Setting DC component of the solution (the mean) to be zero. This is also
    # necessary because the source term to the Poisson equation has zero mean
    # and so the DC component comes out to be ∞.
    ϕ[1, 1, 1] = 0

    solver.transforms.IFFTxy! * ϕ  # Calculate IFFTˣʸ(ϕ) in place.
    solver.transforms.IDCTz!  * ϕ  # Calculate IDCTᶻ(ϕ) in place.

    # Must normalize by 2N for each dimension transformed via FFTW.REDFT.
    @. ϕ = ϕ / (2Nz)

    return nothing
end

#####
##### GPU horizontally periodic pressure solver
#####

function HorizontallyPeriodicPressureSolver(::GPU, grid, pressure_bcs, no_args...)
    Nx, Ny, Nz, _ = unpack_grid(grid)

    kx², ky², kz² = generate_discrete_eigenvalues(grid, pressure_bcs) .|> CuArray
    wavenumbers = (kx² = kx², ky² = ky², kz² = kz²)

    kz⁺ = reshape(0:Nz-1,       1, 1, Nz)
    kz⁻ = reshape(0:-1:-(Nz-1), 1, 1, Nz)

    ω_4Nz⁺ = ω.(4Nz, kz⁺) |> CuArray
    ω_4Nz⁻ = ω.(4Nz, kz⁻) |> CuArray

    # The zeroth coefficient of the IDCT (DCT-III or FFTW.REDFT01) is not
    # multiplied by 2. For some reason, we only need to account for this when
    # doing a 1D IDCT (for PPN boundary conditions) but not for a 2D IDCT. It's
    # possible that the masks are effectively doing this job.
    ω_4Nz⁻[1] *= 1/2

    constants = (ω_4Nz⁺ = ω_4Nz⁺, ω_4Nz⁻ = ω_4Nz⁻)

    # Storage for RHS and Fourier coefficients is hard-coded to be Float64
    # because of precision issues with Float32.
    # See https://github.com/climate-machine/Oceananigans.jl/issues/55
    storage = zeros(Complex{Float64}, Nx, Ny, Nz) |> CuArray

    @debug "Planning transforms for PressureSolver{HorizontallyPeriodic, GPU}..."
    x_bc, y_bc, z_bc = pressure_bcs.x.left, pressure_bcs.y.left, pressure_bcs.z.left
    FFTxy!  = plan_forward_transform(storage, x_bc, [1, 2])
    FFTz!   = plan_forward_transform(storage, z_bc, 3)
    IFFTxy! = plan_backward_transform(storage, x_bc, [1, 2])
    IFFTz!  = plan_backward_transform(storage, z_bc, 3)
    @debug "Planning transforms for PressureSolver{HorizontallyPeriodic, GPU} done!"

    transforms = ( FFTxy! =  FFTxy!,  FFTz! =  FFTz!,
                  IFFTxy! = IFFTxy!, IFFTz! = IFFTz!)

    return PressureSolver(HorizontallyPeriodic(), GPU(), wavenumbers, storage, transforms, constants)
end

"""
    solve_poisson_equation!(solver::PressureSolver{HorizontallyPeriodic, GPU}, grid)

Similar to solve_poisson_equation!(solver::HPPS{CPU}, grid) except that since
the discrete cosine transform is not available through cuFFT, we perform our
own fast cosine transform (FCT) via an algorithm that utilizes the FFT.

Note that for the FCT algorithm to work, the input must have been permuted along
the dimension the FCT is to be calculated by ordering the odd elements first
followed by the even elements. For example,

    [a, b, c, d, e, f, g, h] -> [a, c, e, g, h, f, d, b]

The output will be permuted in this way and so the permutation must be undone.

We should describe the algorithm in detail in the documentation.
"""
function solve_poisson_equation!(solver::PressureSolver{HorizontallyPeriodic, GPU}, grid)
    kx², ky², kz² = solver.wavenumbers.kx², solver.wavenumbers.ky², solver.wavenumbers.kz²
    ω_4Nz⁺, ω_4Nz⁻ = solver.constants.ω_4Nz⁺, solver.constants.ω_4Nz⁻

    # We can use the same storage for the RHS and the solution ϕ.
    RHS, ϕ = solver.storage, solver.storage

    # Calculate DCTᶻ(f) in place using the FFT.
    solver.transforms.FFTz! * RHS
    @. RHS = 2 * real(ω_4Nz⁺ * RHS)

    solver.transforms.FFTxy! * RHS  # Calculate FFTˣʸ(f) in place.

    # Solve the discrete Poisson equation in spectral space. We are essentially
    # computing the Fourier coefficients of the solution from the Fourier
    # coefficients of the RHS.
    @. ϕ = -RHS / (kx² + ky² + kz²)

    # Setting DC component of the solution (the mean) to be zero. This is also
    # necessary because the source term to the Poisson equation has zero mean
    # and so the DC component comes out to be ∞.
    ϕ[1, 1, 1] = 0

    solver.transforms.IFFTxy! * ϕ  # Calculate IFFTˣʸ(ϕ̂) in place.

    # Calculate IDCTᶻ(ϕ̂) in place using the FFT.
    ϕ .*= ω_4Nz⁻
    solver.transforms.IFFTz! * ϕ

    return nothing
end


text:
@testset "Initial population plots" begin
    mpSim2 = MPsim()
    filePath = joinpath( "base", "baseConfig" )
    fileName = string( filePath, ".xlsx" )

    XLSX.openxlsx( fileName ) do xf
        MPIO.readInitPop( mpSim2, xf, filePath )
    end  # XLSX.openxlsx( fileName ) do xf

    initialPopPlot( mpSim2 )
    initialPopAgePlot( mpSim2, 12, :age )
    initialPopAgePlot( mpSim2, 12, :tenure )
    initialPopAgePlot( mpSim2, 12, :timeInNode )
    #=
    initialPopPlot( mpSim2, showplot=false, filename="plot/initpop" )
    initialPopAgePlot( mpSim2, 12, :age, showplot=false,
        filename="plot/initpop age" )
    initialPopAgePlot( mpSim2, 12, :tenure, showplot=false,
        filename="plot/initpop tenure" )
    initialPopAgePlot( mpSim2, 12, :timeInNode, showplot=false,
        filename="plot/initpop time in node" )
    =#
end  # @testset "Initial population plots"


text:
#=
NOTE:
if allowing star like bond, the graph is not simple
also, a tensor network can be a multi-graph, it can not be avoided in the contraction process
note: not all legs form bonds
note: single site bond means trace
note: with star contraction, the relation between elimination ordering and contraction ordering still holds
note: general graph means the equivalence between edges and vertices
=#

export IncidenceMatrix, imap
export linegraph_imap
export nv, ne, vertices, edges, count_vertices, neighbors, findedges
export add_vertex, rem_edge, rem_vertex, add_vertex!, rem_edge!, rem_vertex!
export isvoid, isloop, isbond, isstar, isdirected_graph, issimple_graph, isconnected
export NotSimpleGraphError

abstract type IncidenceMatrix end

########### interfaces ##############
"""
    imap(g::IncidenceMatrix) -> AbstractArray

get the incidence matrix"""
function imap end

"""remove an edge"""
function rem_edge! end

"""remove a vertex"""
function rem_vertex! end

"""add a new vertex"""
function add_vertex! end

Base.show(io::IO, g::IncidenceMatrix) = show(io, "plain/text", g)
function Base.show(io::IO, ::MIME"plain/text", g::IncidenceMatrix)
    print(io, "$(typeof(g).name)")
    print(io, "(V$(nv(g)),E$(ne(g))) ")
    print(io, join(["f$ie($(join(vertices(g, ie), ",")))" for ie in edges(g)], " * "))
end

# NOTE: `copy` in base is also required.

############ derived fallback functions with above interfaces ##############
for FUNC in [:rem_edge, :rem_vertex, :add_vertex]
    IFUNC = Symbol(FUNC, :!)
    @eval $FUNC(fg::IncidenceMatrix, args...) = $IFUNC(copy(fg), args...)
    @eval $FUNC(args...) = fg::IncidenceMatrix -> $FUNC(fg, args...)
    @eval $IFUNC(args...) = fg::IncidenceMatrix -> $IFUNC(fg, args...)
end

nv(fg::IncidenceMatrix) = size(fg |> imap, 1)
ne(fg::IncidenceMatrix) = size(fg |> imap, 2)

vertices(fg::IncidenceMatrix) = 1:nv(fg)
vertices(fg::IncidenceMatrix, ie::Int) = findall(!iszero, @inbounds view(fg |> imap, :, ie))
edges(fg::IncidenceMatrix) = 1:ne(fg)
edges(fg::IncidenceMatrix, it::Int) = findall(!iszero, @inbounds view(fg |> imap, it, :))

count_vertices(fg::IncidenceMatrix, ie::Int) = count(x->x!=0, view(fg |> imap, :,ie))
for FUNC in [:count_vertices]
    @eval $FUNC(fg::IncidenceMatrix) = $FUNC.(Ref(fg), fg |> edges)
end

# boolean functions about local features
isvoid(fg::IncidenceMatrix, ie::Int) = count_vertices(fg, ie) == 0
isloop(fg::IncidenceMatrix, ie::Int) = count_vertices(fg, ie) == 1
isbond(fg::IncidenceMatrix, ie::Int) = count_vertices(fg, ie) == 2
isstar(fg::IncidenceMatrix, ie::Int) = count_vertices(fg, ie) > 2
"""
    isconnected(fg::IncidenceMatrix, vertices, [edge]) -> Bool

Return true if vertices are connected (by edge).
"""
isconnected(fg::IncidenceMatrix, vertices, edge) = all(vertex->imap(fg)[vertex, edge] .!= 0, vertices)
isconnected(fg::IncidenceMatrix, vertices) = any(edge->isconnected(fg, vertices, edge), edges(fg))

# boolean function about global features
isdirected_graph(fg::IncidenceMatrix) = false
issimple_graph(fg::IncidenceMatrix) = all(ie->isbond(fg, ie), fg |> edges)

"""
    neighbors(fg::IncidenceMatrix, it) -> Vector

the neighbors of a vertex.
"""
function neighbors(fg::IncidenceMatrix, it)
    egs = vcat(edges.(Ref(fg), it)...)
    nbs = Int[]
    for iedge in egs
        for jt in vertices(fg, iedge)
            if !(jt in it) && !(jt in nbs)
                push!(nbs, jt)
            end
        end
    end
    nbs
end

linegraph_imap(fg::IncidenceMatrix) = imap(fg)'

"""
    findedges([func], fg::IncidenceMatrix, vertices) -> edge(s)

find edges relating vertices, alternative argument `func` is the finder, it can be `findfirst`, `findall`...
"""
findedges(func, fg::IncidenceMatrix, vertices) = func(edge->isconnected(fg, vertices, edge), edges(fg))
findedges(fg::IncidenceMatrix, vertices) = findedges(findall, fg, vertices)

######### error handling and warning  ############
struct NotSimpleGraphError <: Exception
    msg::String
end


text:
# Use baremodule to shave off a few KB from the serialized `.ji` file
baremodule libvorbis_jll
using Base
using Base: UUID
import JLLWrappers

JLLWrappers.@generate_main_file_header("libvorbis")
JLLWrappers.@generate_main_file("libvorbis", UUID("f27f6e37-5d2b-51aa-960f-b287f2bc3b7a"))
end  # module libvorbis_jll


text:
# VSCode users:
# Do not run next line of code with Ctrl-Enter!
# Run next line of code with Alt-Enter (run line), Shift-Enter (run cell), or Ctrl-F5 (run file)
mypath = dirname(@__FILE__)

## ##########################################################################
# Initialization

using ModelBaseEcon
using StateSpaceEcon
using TimeSeriesEcon

using Test
using Plots
using Random
using Distributions

# Fix the random see for reproducibility
Random.seed!(1234);

nothing
## ##########################################################################

# ##  The Model File

# We recommend placing the model definition in its own Julia module in a separate
# source file. Although this is not strictly necessary, it helps to keep the code
# well organized and it also allows us to take advantage of pre-compilation. The
# first time we load the model file it takes some time to compile, and after that
# loading is much faster.

# The FRB/US model we will be working with is located in models/FRBUS_VAR.jl.
# This file was automatically generated from the `model.xml` file contained
# within frbus_package.zip.

# ### Some Notes About the Julia Model File

# 1. Some variables are declared as *log variables* using the `@log` declaration
#    within a `@variables` block. For example

#    ```julia
#    @variables model begin
#        # ... #
#        "Investment in equipment, current \$" @log ebfin
#        "Personal consumption expenditures, current \$ (NIPA definition)" @log ecnian
#        # ... #
#    end
#    ```

#     A full discussion of log variables is beyond the scope of this tutorial.
#     However a very simplified explanation is that this improves the stability of
#     the numerical solver for variables which are always positive.

# 2. Variables which do not have an associated equation, and for which data is
#    always given, are declared *exogenous* using an `@exogenous` block. For
#    example

#    ```julia
#    @exogenous model begin
#        # Exogenous variables:
#        "Potential government employment ratio (relative to business)" adjlegrt
#        "Dummy, post-1979 indicator" d79a
#        "Dummy, 1980-1995 indicator" d8095
#        # ... #
#    end
#    ```

# 3. The EViews syntax is translated to Julia syntax. EViews functions `d()` and
#    `dlog()` are replaced with their equivalent StateSpaceEcon *meta functions*
#    `@d()` and `@dlog()`. EViews `@movav()` is left alone, because a *meta
#    function* by the same name already exists in StateSpaceEcon, and does the
#    same thing. Finally, the EViews `@recode()` is replaced with the equivalent
#    Julia function `ifelse()`, or where appropriate with a `min()` or a `max()`.

# 4. Several equations contain expressions matching the pattern
#    `1 / (1 + exp(-cx))`, where `c` is a large constant (usually 25) and `x` is
#    some expression. This sigmoid function is a smooth approximation of the
#    Heaviside step function for large values of `c`. While mathematically the
#    derivative of this function is well defined and converges to approximately
#    zero everywhere outside a very small interval containing 0, numerically it
#    causes problems because it results in either 0/0 or ∞/∞. To remedy this
#    situation, we replace such patterns with equivalent calls to `heaviside(cx)`,
#    where the function `heaviside()` is defined in the model file.

#    ```julia
#    export heaviside
#    "Heaviside step function" @inline heaviside(x) = convert(typeof(x), x>zero(x))
#    ```

# ### Regenerating the Model File

# We have included the script update_models.jl. It is not necessary for
# running the code below, but it may be helpful for further experimentation.

# If the model file is missing, for some reason or another, this script will
# automatically download `frbus_package.zip` and process the `model.xml` within it
# to re-generate `model/FRBUS_VAR.jl`. This could also be useful if you make
# modifications to `model.xml` (including not only the equations, but also the
# parameter values), or if you want to use a different FRBUS package from the one
# posted on the FRBUS website. In this case, simply place your `frbus_package.zip`
# in the `models/` directory and run `update_models.jl`. Of course such
# modifications can also be done directly into the existing `models/FRBUS_VAR.jl`
# file.

# Note that after updating `models/FRBUS_VAR.jl`, it'd be best to restart the
# REPL. The first time you load the new model module it'll take a bit longer due
# to pre-compilation.

# Update models files if necessary
models_path = joinpath(mypath, "models")

if !isfile(joinpath(models_path, "FRBUS_VAR.jl")) ||
    (isfile(joinpath(models_path, "frbus_package.zip")) &&
    mtime(joinpath(models_path, "FRBUS_VAR.jl")) < mtime(joinpath(models_path, "frbus_package.zip")))
    @info "Updating model files"
    include("update_models.jl")
end

## ##########################################################################
# ## Load the Model

# Assuming that the `models/` directory is already in the `LOAD_PATH` list, we can
# load the model by `using` its module. Once loaded, the module contains a
# variable `model` which represents the model object.

unique!(push!(LOAD_PATH, joinpath(mypath, "models")))
using FRBUS_VAR
m = FRBUS_VAR.model

## ###########
# We see that the model has a number of variables, shocks, equations, and
# parameters. The total number of variables include exogenous variables.

m.variables

## #######
# We also see that the model object includes a number of auxiliary equations.
# These equations (and variables) are automatically added as substitutions for
# expressions that must be positive.

m.auxeqns

## #######
# For example, we see that variable `aux1` was added with the first equation in
# the above list. At the same time, in equation for `dpgap`, the expression
# `log(phr[t] * pxp[t])` has been replaced by `aux1[t]`.

m.equations[2]

# A detailed discussion of auxiliary variables and equations is beyond the scope
# of this tutorial. It suffices to say that we can safely ignore their presence
# for now.

## ##########################################################################
# ## Load the Longbase Data

# Unfortunately the `longbase` data is available only in EViews format, which
# cannot be read automatically by open source software (at least to our
# knowledge). For convenience, here we have included the version of `longbase` from
# 23-07-2020 in a csv format and a function that loads that data. The function is
# defined in file load_longbase.jl. Note that this is not a
# module, so we load it by calling `include()`, not `using`.

include("load_longbase.jl")
longbase = load_longbase("longbase_23072020.csv")

## ##########################################################################
# ## Load `set_policy.jl`

# The model contains a number of switch variables to control which monetary policy
# function is used and which fiscal policy function is used at each period of the
# simulation. For convenience, we have included functions `set_mp!()` and
# `set_fp!()`, which are defined in set_policy.jl.

# load functions set_mp! and set_fp!
include("set_policy.jl")

# check out the docs. In the REPL you can also type ?set_mp!
@doc set_mp!

# these are the mp switches
dmp_switches

# and the docs for setting fiscal policy
@doc set_fp!

# these are the fp switches
dfp_switches

## ##########################################################################
# ## Prepare the Simulation Plan

# The simulation is controlled by a `Plan` object. The plan is defined by
# a model object and a simulation range. The full range handled by the plan
# contains additional periods before and after the simulation range, which account
# for initial and final conditions. By default, the simulation plan is setup such
# that all shocks are exogenous and all variables are endogenous, except for the
# variables that are declared either in an `@exogenous` block or with the `@exog`
# declaration within an `@variables` block in the model file.

sim = 2020Q1:2025Q4     # simulation range
p = Plan(m, sim)        # the plan object

ini = firstdate(p):first(sim) - 1      # range of initial conditions
fin = last(sim) + 1:lastdate(p)        # range of final conditions

# Note that the `fin` range is actually empty.
# This is because this model doesn't have any leads.
isempty(fin)

## ##########################################################################
# ## Prepare the Exogenous Data

# We start by pre-allocating simulation data that is set to 0 everywhere.  Then we
# assign within it the data from `longbase` using Julia's `.=` operator.

# ed = exogenous data
ed = zerodata(m, p);
ed .= longbase[p.range];

# Next we set the monetary policy, the fiscal policy and a few other switches.

# set monetary policy
set_mp!(ed, :dmpintay);

# turn off zero bound and policy thresholds;
# hold policy maker's perceived equilibrium real interest rate
ed.dmptrsh .= 0.0;
ed.rffmin .= -9999;
ed.drstar .= 0.0;

# set fiscal policy
set_fp!(ed, :dfpsrp);

## ##########################################################################
# ## Back out the Shocks

# The first simulation test is to compute the shocks given the variable paths from
# `longbase`. To do this, we swap the variables and shocks, making variables
# exogenous and shocks endogenous. The mapping between variables and their
# corresponding shocks is declared in the model file, so we can simply call
# `autoexogenize!`. We make a copy of the plan `p`, so that the original
# plan would not be modified. We also make a copy of the exogenous data, `ed`, so
# that the original would remain unchanged.

p_0 = autoexogenize!(copy(p), m, sim)
ed_0 = copy(ed)

# Now we run the `simulate` command. Note that the first time we run a
# function in Julia it takes a bit longer due to compilation time. In this case,
# it takes much longer because the model is very large and each and every equation
# gets compiled together with its automatic derivative.

sol_0 = @time simulate(m, ed_0, p_0; verbose=true, tol=1e-12);

# The compilation is done once and the compiled code is used in every call after
# that. So the second call to `simulate` is much, much faster.

sol_0 = @time simulate(m, ed_0, p_0; verbose=true, tol=1e-12);

## ##########################################################################
# ## Recover the Baseline Case

# Next simulation is a sanity check test. If we run a simulation with the shocks
# set to the values we just backed out, the resulting variable paths must match
# the ones we started with.

# Once again we start with an exogenous data set everywhere to 0. Then we assign
# only the initial conditions and the shocks we just backed out.

p_r = Plan(m, sim);
ed_r = zerodata(m, p_r);

# initial conditions for the variables are taken from longbase
ed_r[ini, m.variables] = longbase[ini, m.variables];

# shocks are taken from from sol_0
ed_r[p_r.range, m.shocks] = sol_0[p_r.range, m.shocks];

# exogenous variables are also taken  from sol_0
exogenous = m.variables[isexog.(m.variables)];
ed_r[p_r.range, exogenous] = sol_0[p_r.range, exogenous];

# Now, the only thing left is to set the initial guess for the endogenous
# variables. If we leave it at 0, that would be an initial guess too far from the
# solution and the Newton-Raphson will likely diverge. If we set it to the known
# solution, that would diminish this exercise to merely verifying that it is
# indeed a solution (we already know that). So, to make things a bit more
# interesting, we add a bit of noise to the true solution.

endogenous = m.variables[.!isexog.(m.variables)];
ed_r[sim, endogenous] = longbase[sim, endogenous] .+ 0.03 .* randn(length(sim), length(endogenous));

# Once again we have to set the monetary policy and the fiscal policy rules, as
# well as the values of some of the other switches.

set_mp!(ed_r, :dmpintay);
ed_r.dmptrsh .= 0.0;
ed_r.rffmin .= -9999;
ed_r.drstar .= 0.0;
set_fp!(ed_r, :dfpsrp);

# And finally we can run the simulation and check to make sure that indeed the
# recovered simulation matches the base case.

sol_r = @time simulate(m, ed_r, p_r, verbose=true, tol=1e-9);
@test sol_r ≈ sol_0

## ##########################################################################
# ## Simulate a shock

# The last exercise is to simulate the impulse response to a unit shock in
# `rffintay`.
m.rffintay

# We start with the base case and add 1 to the `rffintay_a` shock at the first
# period of the simulation.

p_1 = Plan(m, sim);
ed_1 = copy(sol_0);

ed_1.rffintay_a[first(sim)] += 1;
sol_1 = @time simulate(m, ed_1, p_1;  verbose=true, tol=1e-9);

## ##########################################################################
# Finally, we can plot the impulse response function to see what we've done.
# compute the differences between the base case and the shocked simulation.
dd = hcat(SimData(p.range),
    d_rff=sol_1.rff - sol_0.rff,
    d_rg10=sol_1.rg10 - sol_0.rg10,
    d_lur=sol_1.lur - sol_0.lur,
    d_pic4=sol_1.pic4 - sol_0.pic4,
);

# produce the plot
plot(dd[sim], vars=(:d_rff, :d_rg10, :d_lur, :d_pic4),
     legend=false, size=(600, 400))


text:
# Plot basic Gaussian Process

using PyPlot, GaussianProcesses

srand(13579)
# Training data
n=10                 #number of training points
x = 2π * rand(n)
y = sin(x) + 0.05*randn(n)

#Select mean and covariance function
mZero = MeanZero()                   #Zero mean function
kern = SE(0.0,0.0)                   #Sqaured exponential kernel (note that hyperparameters are on the log scale)

gp = GP(x,y,mZero,kern,-1.0)      #Fit the GP, where -1.0 is the log Gaussian noise
plot(gp)                          #Plot the GP

optimize!(gp; method=Optim.BFGS())   #Optimise the hyperparameters

plot(gp)   #Plot the GP after the hyperparameters have been optimised

# Add observation
push!(gp, [1.0], [2.0])


text:
# using Revise
using PiecewiseDeterministicMarkovProcesses, LinearAlgebra, Random, DifferentialEquations, Sundials
const r = 10.

function AnalyticalSample(xc0,xd0,ti,nj::Int64)
	xch = [xc0[1]]
	xdh = [xd0[1]]
	th  = [ti]
	list_rng = Float64[]
	t = ti
	while length(th)<nj
		xc = xch[end]
		xd = xdh[end]
		push!(list_rng, rand())
		S = -log(list_rng[end])
		a = -r * (2mod(xd,2)-1)
		dt = log(a*S/xc+1)/a
		t += dt
		push!(th, t)
		push!(xch,xc + a * S )
		push!(xdh,xd .+ 1 )
		push!(list_rng, rand())
		S = -log(list_rng[end])
	end
	return th, xch, xdh, list_rng
end

function F!(ẋ, xc, xd, parms, t)
	ẋ[1] = -r * (2mod(xd[1],2)-1) * xc[1]
end

R(x) = x

function R!(rate, xc, xd, parms, t, issum::Bool)
	# rate fonction
	if issum==false
		rate[1] = R(xc[1])
		rate[2] = 0.0
		return R(xc[1]),40.
	else
		return R(xc[1]),40.
	end
end

xc0 = [1.0]
xd0 = [0,0]

nu = [[1 0];[0 -1]]
parms = [0.0]
ti = 0.332
tf = 100000.
nj = 50

Random.seed!(8)
	res_a = AnalyticalSample(xc0,xd0,ti,nj)

errors = Float64[]
# state of the random generator
rnd_state = 0.

println("\n\nComparison of solvers")
	for ode in [(:lsoda,"lsoda"),(:cvode,"cvode"),(CVODE_BDF(),"CVODEBDF"),(CVODE_Adams(),"CVODEAdams"),(Tsit5(),"tsit5"),(Rodas4P(autodiff=false),"rodas4P-noAutoDiff"),(Rodas5(),"rodas5"),(Rosenbrock23(),"RS23"),(AutoTsit5(Rosenbrock23()),"AutoTsit5-RS23")]
		Random.seed!(8)

		problem = PDMP.PDMPProblem(F!, R!, nu, xc0, xd0, parms, (ti, tf))
		res =  PDMP.solve(problem, CHV(ode[1]); n_jumps = nj)
		# this is to check the state of the random generator at the end of the simulation
		if ode[1] == :lsoda
			global rnd_state = rand()
		else
			@assert rnd_state == rand()
		end

		println("--> norm difference = ", norm(res.time - res_a[1],Inf64), "  - solver = ",ode[2])
		push!(errors,norm(res.time - res_a[1],Inf64))
	end


text:
# compiler support for working with run-time libraries

function link_library!(job::CompilerJob, mod::LLVM.Module, lib::LLVM.Module)
    # linking is destructive, so copy the library
    lib = LLVM.Module(lib)

    # save list of external functions
    exports = String[]
    for f in functions(mod)
        fn = LLVM.name(f)
        if !haskey(functions(lib), fn)
            push!(exports, fn)
        end
    end

    link!(mod, lib)

    ModulePassManager() do pm
        # internalize all functions that aren't exports
        internalize!(pm, exports)

        # eliminate all unused internal functions
        global_optimizer!(pm)
        global_dce!(pm)
        strip_dead_prototypes!(pm)

        run!(pm, mod)
    end
end

const libcache = Dict{String, LLVM.Module}()

#
# CUDA device library
#

function get_libdevice(cap)
    if isa(libdevice[], Dict)
        # select the most recent & compatible library
        vers = keys(libdevice[])
        compat_vers = Set(ver for ver in vers if ver <= cap)
        isempty(compat_vers) && error("No compatible CUDA device library available")
        ver = maximum(compat_vers)
        path = libdevice[][ver]
    else
        libdevice[]
    end
end

function load_libdevice(cap)
    path = get_libdevice(cap)

    get!(libcache, path) do
        open(path) do io
            parse(LLVM.Module, read(path), JuliaContext())
        end
    end
end

function link_libdevice!(job::CompilerJob, mod::LLVM.Module, lib::LLVM.Module)
    # override libdevice's triple and datalayout to avoid warnings
    triple!(lib, triple(mod))
    datalayout!(lib, datalayout(mod))

    link_library!(job, mod, lib)

    ModulePassManager() do pm
        push!(metadata(mod), "nvvm-reflect-ftz",
              MDNode([ConstantInt(Int32(1), JuliaContext())]))
        # TODO: run the reflect pass?
        run!(pm, mod)
    end
end

#
# CUDAnative run-time library
#

# get the path to a directory where we can put cache files (machine-specific, ephemeral)
# NOTE: maybe we should use XDG_CACHE_PATH/%LOCALAPPDATA%, but other Julia cache files
#       are put in .julia anyway so let's just follow suit for now.
function cachedir(depot=DEPOT_PATH[1])
    # this mimicks Base.compilecache. we can't just call the function, or we might actually
    # _generate_ a cache file, e.g., when running with `--compiled-modules=no`.
    if VERSION >= v"1.3.0-alpha.146"
        entrypath, entryfile = Base.cache_file_entry(Base.PkgId(CUDAnative))
        abspath(depot, entrypath, entryfile)
    else
        cachefile = abspath(depot, Base.cache_file_entry(Base.PkgId(CUDAnative)))

        # the cachefile consists of `/depot/compiled/vXXX/CUDAnative/$slug.ji`
        # transform that into `/depot/compiled/vXXX/CUDAnative/$slug/`
        splitext(cachefile)[1]
    end
end

# remove the existing cache globally, so any change to CUDAnative triggers recompilation.
rm(cachedir(); recursive=true, force=true)
# create an empty cache directory. since we only ever load from the first existing cachedir,
# this effectively invalidates preexisting caches in lower layers of the depot.
mkdir(cachedir())

## higher-level functionality to work with runtime functions

function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[])
    bb = position(builder)
    f = LLVM.parent(bb)
    mod = LLVM.parent(f)

    # get or create a function prototype
    if haskey(functions(mod), rt.llvm_name)
        f = functions(mod)[rt.llvm_name]
        ft = eltype(llvmtype(f))
    else
        ft = LLVM.FunctionType(rt.llvm_return_type, rt.llvm_types)
        f = LLVM.Function(mod, rt.llvm_name, ft)
    end

    # runtime functions are written in Julia, while we're calling from LLVM,
    # this often results in argument type mismatches. try to fix some here.
    for (i,arg) in enumerate(args)
        if llvmtype(arg) != parameters(ft)[i]
            if (llvmtype(arg) isa LLVM.PointerType) &&
               (parameters(ft)[i] isa LLVM.IntegerType)
                # Julia pointers are passed as integers
                args[i] = ptrtoint!(builder, args[i], parameters(ft)[i])
            else
                error("Don't know how to convert ", arg, " argument to ", parameters(ft)[i])
            end
        end
    end

    call!(builder, f, args)
end

## functionality to build the runtime library

function emit_function!(mod, cap, f, types, name)
    tt = Base.to_tuple_type(types)
    new_mod, entry = codegen(:llvm, CompilerJob(f, tt, cap, #=kernel=# false);
                             libraries=false, strict=false)
    LLVM.name!(entry, name)
    link!(mod, new_mod)
end

function build_runtime(cap)
    mod = LLVM.Module("CUDAnative run-time library", JuliaContext())

    for method in values(Runtime.methods)
        emit_function!(mod, cap, method.def, method.types, method.llvm_name)
    end

    mod
end

function load_runtime(cap)
    # find the first existing cache directory (for when dealing with layered depots)
    cachedirs = [cachedir(depot) for depot in DEPOT_PATH]
    filter!(isdir, cachedirs)
    @assert !isempty(cachedirs)
    input_dir = first(cachedirs)

    # we are only guaranteed to be able to write in the current depot
    output_dir = cachedir()

    # if both aren't equal, copy pregenerated runtime libraries to our depot
    # NOTE: we don't just lazily read from the one and write to the other, because
    #       once we generate additional runtimes in the output dir we don't know if
    #       it's safe to load from other layers (since those could have been invalidated)
    if input_dir != output_dir
        mkpath(dirname(output_dir))
        cp(input_dir, output_dir)
    end

    name = "runtime_$(cap.major)$(cap.minor).bc"
    path = joinpath(output_dir, name)

    get!(libcache, path) do
        if ispath(path)
            open(path) do io
                parse(LLVM.Module, read(io), JuliaContext())
            end
        else
            @info "Building the CUDAnative run-time library for your sm_$(cap.major)$(cap.minor) device, this might take a while..."
            lib = build_runtime(cap)
            open(path, "w") do io
                write(io, lib)
            end
            lib
        end
    end
end