text:
"""
    struct ElementIterator{E,M}

Return an iterator for iterating over all elements of type `E` on objects of
type `M`.
"""
struct ElementIterator{E,M} <: AbstractVector{E}
    mesh::M
end

ElementIterator{E}(mesh::M) where {E,M <: AbstractMesh} = ElementIterator{E,M}(mesh)
ElementIterator(mesh,E) = ElementIterator{E}(mesh)

mesh(iter::ElementIterator) = iter.mesh

function Base.getindex(iter::ElementIterator,I)
    @assert eltype(I) <: Integer
    map(i->iter[i],I)
end

function Base.size(iter::ElementIterator{<:LagrangeElement,<:GenericMesh})
    msh               = mesh(iter)
    E                 = eltype(iter)
    tags::Matrix{Int} = msh.elements[E]
    Np, Nel           = size(tags)
    return (Nel,)
end

function Base.getindex(iter::ElementIterator{<:LagrangeElement,<:GenericMesh},i::Int)
    E                   = eltype(iter)
    mesh                = iter.mesh
    tags::Matrix{Int}   = mesh.elements[E]
    node_tags           = view(tags,:,i)
    vtx                 = view(mesh.nodes,node_tags)
    el                  = E(vtx)
    return el
end

function Base.size(iter::ElementIterator{<:ParametricElement,<:GenericMesh})
    E               = eltype(iter)
    mesh            = iter.mesh
    tags::Vector{E} = mesh.elements[E]
    return (length(iter.mesh.elements[E]),)
end

function Base.getindex(iter::ElementIterator{<:ParametricElement,<:GenericMesh},i::Int)
    E               = eltype(iter)
    mesh            = iter.mesh
    els::Vector{E} = mesh.elements[E]
    return els[i]
end

# iterator for submesh
function Base.size(iter::ElementIterator{<:AbstractElement,<:SubMesh})
    E          = eltype(iter)
    submesh    = iter.mesh
    idxs       = dom2elt(submesh,E)
    return (length(idxs),)
end

function Base.getindex(iter::ElementIterator{<:AbstractElement,<:SubMesh},i::Int)
    E      = eltype(iter)
    msh    = mesh(iter) # a SubMesh
    p_msh  = parent(msh) # parent mesh
    iglob  = dom2elt(msh,E)[i] # global index of element in parent mesh
    iter = ElementIterator(p_msh,E)
    return iter[iglob]
end

ElementIterator(m::CartesianMesh) = ElementIterator(m,etype(m))

# element iterator for cartesian mesh
# FIXME: the ElementIterator for a CartesianMesh should inherit from
# AbstractMatrix and not AbstractVector since these are more naturally indexed
# as matrices. Also this should be made generic on the dimension
function Base.size(iter::ElementIterator{<:Any,<:CartesianMesh})
    E       = eltype(iter)
    grids   = grid1d(iter.mesh)
    sz      = size(iter.mesh) .- 1
    return (prod(sz),)
end

# 1d case
function Base.getindex(iter::ElementIterator{<:Any,<:CartesianMesh{1}}, i::Int)
    E           = eltype(iter)
    xx          = xgrid(iter.mesh)
    low_corner  = (xx[i], )
    high_corner = (xx[i+1], )
    el          = E(low_corner,high_corner) # construct the element
    return el
end

# 2d case
function Base.getindex(iter::ElementIterator{<:Any,<:CartesianMesh{2}}, i::Int,j::Int)
    E      = eltype(iter)
    xx     = xgrid(iter.mesh)
    yy     = ygrid(iter.mesh)
    low_corner = (xx[i],yy[j])
    high_corner = (xx[i+1],yy[j+1])
    el  = E(low_corner,high_corner) # construct the element
    return el
end

function Base.getindex(iter::ElementIterator{<:Any,<:CartesianMesh{2}}, I::Int)
    sz  = size(iter.mesh) .- 1
    idx = CartesianIndices(sz)[I]
    iter[idx[1],idx[2]]
end

function Base.iterate(iter::ElementIterator{<:Any,<:CartesianMesh},state=1)
    n = length(iter)
    if state > n
        return nothing
    else
        iter[state],state+1
    end
end

max_stars_repo_name: WaveProp/WaveProp

max_issues_repo_name: WaveProp/WaveProp

max_forks_repo_name: WaveProp/WaveProp


text:
# --- PRIVATE METHODS --------------------------------------------------------------------------------------- #
function _build_call_contract_object(data::Dict{String,Any})

    # grab stuff from data -
    # "symbol":"AMD",
    # "type":"call",
    # "sense":"sell",
    # "strike_price":60.0,
    # "premium_value":1.50,
    # "number_of_contracts":1,
    # "expiration":"2020-08-07"

    symbol = data["ticker_symbol"]
    sense = data["sense"]
    strike_price = data["strike_price"]
    premium_value = data["premium_value"]
    number_of_contracts = data["number_of_contracts"]
    expiration_date_string = data["expiration"]
    contract_multuplier = data["contract_multiplier"]

    # TODO - check sense, is this legit?
    # TODO - check is date string formatting correct?

    # build -
    callOptionContract = PSCallOptionContract(symbol, Date(expiration_date_string), strike_price, premium_value, number_of_contracts;
        sense=Symbol(sense), contractMultiplier=contract_multuplier)

    # return -
    return PSResult(callOptionContract)
end

function _build_put_contract_object(data::Dict{String,Any})

    symbol = data["ticker_symbol"]
    sense = data["sense"]
    strike_price = data["strike_price"]
    premium_value = data["premium_value"]
    number_of_contracts = data["number_of_contracts"]
    expiration_date_string = data["expiration"]
    contract_multiplier = data["contract_multiplier"]

    # TODO - check sense, is this legit?
    # TODO - check is date string formatting correct?

    # build -
    putOptionContract = PSPutOptionContract(symbol, Date(expiration_date_string), strike_price, premium_value, number_of_contracts;
        sense=Symbol(sense),contractMultiplier=contract_multiplier)

    # return -
    return PSResult(putOptionContract)
end

function _build_equity_object(data::Dict{String,Any})

    symbol = data["ticker_symbol"]
    purchase_price = data["purchase_price_per_share"]
    number_of_shares = data["number_of_shares"]
    purchase_date_string = data["purchase_date"]

    # TODO - check sense, is this legit?
    # TODO - check is date string formatting correct?

    # build -
    equityObject = PSEquityAsset(ticker_symbol, purchase_price, number_of_shares, Date(purchase_date_string))

    # return -
    return PSResult(equityObject)
end
# ----------------------------------------------------------------------------------------------------------- #

# --- PUBLIC METHODS ---------------------------------------------------------------------------------------- #
function build_simulation_contract_set(simulation_dictionary::Dict{String,Any})::PSResult

    # initialize -
    asset_set = Set{PSAbstractAsset}()

    # TODO: check - do we have the correct keys?

    # grab the list of asset dictionaries -
    asset_dictionary_array = simulation_dictionary["contract_set_parameters"]
    for (index, asset_dictionary) in enumerate(asset_dictionary_array)

        # initialize -
        local result = nothing

        # ok, so lets grab data from the asset_dictionary, and build each asset type -
        type_string = asset_dictionary["type"]
        type_symbol = Symbol(type_string)
        if (type_symbol == :call)
            result = _build_call_contract_object(asset_dictionary)
        elseif (type_symbol == :put)
            result = _build_put_contract_object(asset_dictionary)
        elseif (type_symbol == :equity)
            result = _build_equity_object(asset_dictionary)
        end

        # grab -
        push!(asset_set, result.value)
    end

    # return -
    return PSResult(asset_set)
end

function build_simulation_contract_set(pathToSimulationFile::String)::PSResult

    # TODO: check - is this a legit path -

    # load the experimet file -
    simulation_dictionary = JSON.parsefile(pathToSimulationFile)

    # build the contract set -
    return build_simulation_contract_set(simulation_dictionary)
end

# -- Binary ---
function build_binary_lattice_data_structure(pathToSimulationFile::String)::PSResult

    # TODO: check - is this a legit path -

    # load the experimet file -
    dictionary = JSON.parsefile(pathToSimulationFile)

    # go -
    return build_binary_lattice_data_structure(dictionary)
end

function build_binary_lattice_data_structure(dictionary::Dict{String,Any})::PSResult

    # check - do we have the correct keys -
    main_key = "binary_lattice_model_parameters"
    if (haskey(dictionary,main_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(main_key) key"))
    end
    lattice_parameters = dictionary[main_key]

    # check the other keys -
    # volatility -
    volatility_key = "volatility"
    if (haskey(lattice_parameters,volatility_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(volatility_key) key"))
    end
    volatility = lattice_parameters[volatility_key]

    # risk free rate -
    risk_free_rate_key = "risk_free_rate"
    if (haskey(lattice_parameters,risk_free_rate_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(risk_free_rate_key) key"))
    end
    risk_free_rate = lattice_parameters[risk_free_rate_key]

    # dividend_rate -
    dividend_rate_key = "dividend_rate"
    if (haskey(lattice_parameters, dividend_rate_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(dividend_rate_key) key"))
    end
    dividend_rate = lattice_parameters[dividend_rate_key]

    # days_to_expiration -
    expiration_key = "days_to_expiration"
    if (haskey(lattice_parameters, expiration_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(expiration_key) key"))
    end
    days_to_expiration = lattice_parameters[expiration_key]

    # build lattice object -
    option_parameters = PSBinaryLatticeModel(volatility, days_to_expiration, risk_free_rate, dividend_rate)

    # return -
    return PSResult(option_parameters)
end

# -- Ternary helper methods -
function build_ternary_lattice_data_structure(pathToSimulationFile::String)::PSResult

    # TODO: check - is this a legit path -

    # load the markey file -
    dictionary = JSON.parsefile(pathToSimulationFile)

    # go -
    return build_ternary_lattice_data_structure(dictionary)
end

function build_ternary_lattice_data_structure(dictionary::Dict{String,Any})::PSResult

    # check - do we have the correct keys -
    main_key = "ternary_lattice_model_parameters"
    if (haskey(dictionary,main_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(main_key) key"))
    end
    lattice_parameters = dictionary[main_key]

    # check the other keys -
    # volatility -
    volatility_key = "volatility"
    if (haskey(lattice_parameters,volatility_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(volatility_key) key"))
    end
    volatility = lattice_parameters[volatility_key]

    # risk free rate -
    risk_free_rate_key = "risk_free_rate"
    if (haskey(lattice_parameters,risk_free_rate_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(risk_free_rate_key) key"))
    end
    risk_free_rate = lattice_parameters[risk_free_rate_key]

    # dividend_rate -
    dividend_rate_key = "dividend_rate"
    if (haskey(lattice_parameters, dividend_rate_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(dividend_rate_key) key"))
    end
    dividend_rate = lattice_parameters[dividend_rate_key]

    # days_to_expiration -
    expiration_key = "days_to_expiration"
    if (haskey(lattice_parameters, expiration_key) == false)
        return PSError(ArgumentError("binary lattice specification is missing the $(expiration_key) key"))
    end
    days_to_expiration = lattice_parameters[expiration_key]

    # build lattice object -
    option_parameters = PSTernaryLatticeModel(volatility, days_to_expiration, risk_free_rate, dividend_rate)

    # return -
    return PSResult(option_parameters)
end

# -- Price range helper methods -
function build_simulation_price_array(pathToSimulationFile::String)::PSResult

    # TODO: check - is this a legit path -

    # load the range file -
    dictionary = JSON.parsefile(pathToSimulationFile)

    # go -
    return build_simulation_price_array(dictionary)
end

function build_simulation_price_array(dictionary::Dict{String,Any})::PSResult

    # grab the asset_price_array parameters -
    asset_price_parameters = dictionary["underlying_price_simulation_range"]
    price_start = asset_price_parameters["underlying_price_start"]
    price_stop = asset_price_parameters["underlying_price_stop"]
    number_of_steps = asset_price_parameters["number_of_steps"]

    # build the array -
    assetPriceArray = collect(range(price_start,stop=price_stop,length=number_of_steps))

    # return -
    return PSResult(assetPriceArray)
end
# ----------------------------------------------------------------------------------------------------------- #

ext: jl


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

using Markdown
using InteractiveUtils

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

# ╔═╡ 416d9c62-d2ef-11eb-3a40-450aee175208
using Fermi.GaussianBasis

# ╔═╡ 71915091-164d-440f-931e-7abaa1764e77
using Plots

# ╔═╡ aec67321-b943-478b-9765-8cb95c0f9d06
using Fermi.Integrals

# ╔═╡ f41f12be-3755-4f73-a183-497ced9afa24
using Fermi.HartreeFock

# ╔═╡ 42b34173-623d-4361-a2ba-b55b0cdfd69f
using Fermi

# ╔═╡ 1ad27f33-fe6c-498a-b4e8-89948dd1f94b
html"""

  Basic structures

 Gustavo Aroeira

Center for Computational Quantum Chemistry

University of Georgia, Athens, GA

United States of America

"""

# ╔═╡ c51d70f8-9065-41f9-aea3-73f7c2aba7e7
md"""
## 1. Goal
The main goal in quantum chemistry is to solve the electronic Schrödinger equation for many-electron systems

$\hat{H}|\Psi\rangle = E|\Psi\rangle$

A first approximation to the fermionic wave function is represented as a Slater determinant

$|\Psi_0\rangle = \mathcal{A}[\phi_1(\mathbf{x_1})\phi_2(\mathbf{x_2})\phi_3(\mathbf{x_3})...\phi_N(\mathbf{x_N})]$

where $\mathcal{A}$ is the antisymmetrizer operator and $\phi_i$ are spin-orbitals, that is, one-electron wave functions used to build the many-electron *ansatz*.

The first task is, therefore, to find optimized one-electron functions (orbitals). For atoms, one can do it numerically in a self-consistent fashion. For molecules however, this becomes untractable. Thus, we introduce the *basis set approximation*, where we represent the orbitals as an expansion of $M$ pre-selected basis functions

$\phi_i = \sum_\mu^M C_{\mu i} \chi_\mu$

This is known as the **L**inear **C**ombination of **A**tomic **O**rbitals (LCAO) as the basis functions are designed to resemble atomic orbitals.
"""

# ╔═╡ dc46fbbd-b96e-4fdd-9300-fa5ca87e86de
md"""
## 2. Basis Functions

Analytical solution to the ground state wave function of the hydrogen atom has the shape of a Slater orbital

$\Psi_\text{H} \sim e^{-r}$

Computing integrals and derivatives of slater type functions is too complicated. Thus, gaussian functions are employed instead. Each basis function in our basis set is a combination of gaussian type functions is to mimicking slater type functions.

$\chi = \sum_i N_i G_i(r,\theta,\phi)$
$G_i(r,\theta,\phi) = Y_{l,m_l}(\theta,\phi) r^l e^{-\zeta_i r^2}$

Each basis function $\chi$ is characterized by three quantities:

1. Angular momentum number $l$
2. A set of contraction coefficients $N_i$
3. A set of exponent values $\zeta_i$
"""

# ╔═╡ 97282505-47ee-4011-8198-bd4a7f8f75bf
md"""
##### Let's build a basis function using Fermi
Import the GaussianBasis module
"""

# ╔═╡ 2f8ca46a-3a92-4d0f-a28b-7e47fc577c80
begin
	l = 1
	exp_vals = [10.0, 5.0, 0.6]
	coef_vals = [0.2, 0.6, 0.7]
	BasisFunction(l, coef_vals, exp_vals)
end

# ╔═╡ 51444ecc-7031-4f9b-bb67-733bb6def9dc
@bind bname html"""Select basis set   <select>
    <optgroup label="Pople">
      <option value="sto-3g">STO-3g</option>
      <option value="sto-6g">STO-6g</option>
	  <option value="6-31g">6-31g</option>
    </optgroup>
    <optgroup label="Dunning">
      <option value="cc-pvdz">cc-pVDZ</option>
      <option value="cc-pvtz">cc-pVTZ</option>
      <option value="cc-pvqz">cc-pVQZ</option>
      <option value="cc-pv5z">cc-pV5Z</option>
    </optgroup>
  </select>"""

# ╔═╡ 3a9d4c13-9d77-4deb-add2-b54e29d59934
begin
	# Bohr radius
	a0 = 0.529177210903

	# Define Slater function
	slater(r) = exp(-r/a0)

	# Gaussian function
	function contracted_gaussian(r,e,c)
	out = zeros(length(r))
	for i = eachindex(e)
		ζ = e[i]
		C = c[i]
		out += C*exp.(-ζ.*r.^2 ./a0^2)
	end
	return out ./ maximum(out)
	end

	# Plot range
	rvals = collect(range(0,2,step=0.01))

	bfp, = GaussianBasis.read_basisset(bname, "H")
	plot(rvals, slater.(rvals), label="Slater orbital")
	plot!(rvals, contracted_gaussian(rvals, bfp.exp, bfp.coef), label=bname)
	xlabel!("Distance from the nucleus (Å)")
	ylabel!("ϕ(r)")
end

# ╔═╡ 20b7f076-7177-41fe-b0d6-abd1b9211360
md"""
## 3. Basis Set
"""

# ╔═╡ 57f7f7c7-8a84-4fd0-b3f9-f49d375bf37c
Atom("H", 1, (0.0,0.0,0.0))

# ╔═╡ 7388eaa1-011c-43df-9986-2f59ba259873
atoms = [Atom("H", 1, (0.0,0.0,0.0)),
	     Atom("H", 1, (0.76,0.0,0.0))]

# ╔═╡ 77186eea-5b85-40c9-91bf-f48707752bb7
mol = Molecule(atoms, 0, 1)

# ╔═╡ ee4dbc82-6258-413c-aca5-1bd1e6ee02fc
begin
	# Create basis functions
	s = BasisFunction(0, [0.5215367271], [0.122])
	p = BasisFunction(1, [1.9584045349], [0.727])

	# Map atoms to basis functions
	shells = Dict(
	atoms[1] => [s],
	atoms[2] => [s,p]
	)
	bset = BasisSet(mol, "Hs/Hsp", shells)
end

# ╔═╡ f17bc090-cad6-427a-ace9-7ed921f8368d
md"""
## 4. Integrals

In order to optimize orbitals several integrals over our basis functions are needed.
For example, the overlp integral

$S_{ij} = \int \phi_i^* \phi_j d\tau = C_{\mu i}^*C_{\nu j}\int \chi_\mu^* \chi_\nu d\tau$

Fermi uses the [libcint](https://github.com/sunqm/libcint) library to compute integrals.
"""

# ╔═╡ 15cfa11b-78e1-456c-82f2-02bcd1f8c6c8
I = IntegralHelper(bset)

# ╔═╡ 3685dd55-cb6a-4e0d-8a64-a76fd29d8e49
I["S"]

# ╔═╡ b6e3d8a9-a125-417c-8af1-145b9050d406
md"""
## 5. Hartree-Fock Orbitals
Having an integral object built, a Hartree--Fock computation can be run to obtain optimized orbitals
"""

# ╔═╡ 6fb41a7c-2ae6-4b98-ba5a-fe122512364a
wfn = RHF(I)

# ╔═╡ 90087552-51d4-42ef-8c3e-e5a61eec4fe1
wfn.orbitals.C

# ╔═╡ 683b958d-7959-484e-baf1-9803094dd994
md"""
## 6. Macros to make life easier

Fermi has a `Options` module that is referred to when some information is not specified. Moreover, users can use the macro `@molecule` and `@set` to change keywords and run computations using the desired settings.
"""

# ╔═╡ 165a0055-0195-42d8-8dd2-a3d07fe67f61
begin
	@molecule {
		H 0.0 0.0 0.0
		H 0.76 0.0 0.0
	}

	@set {
		basis 6-31g
		diis false
	}

	@energy rhf
end

# ╔═╡ 9ab18d9f-28a2-43ae-a9df-012223dea282
begin
	Ecc = []
	Emp2 = []
	Eref = []
	dist = collect(range(0.4,3.0,step=0.1))
	for d in dist
		Fermi.Options.set("molstring",
			"""
			H 0.0 0.0 0.0
			H $d 0.0 0.0
			""")

		Iao = IntegralHelper()
		ref = @energy Iao=>rhf
		Imo = IntegralHelper(orbitals=ref.orbitals)
		cc_wfn = @energy Imo,Iao=>ccsd
		push!(Eref, ref.energy)
		push!(Emp2, cc_wfn.guessenergy)
		push!(Ecc, cc_wfn.energy)
	end
end

# ╔═╡ 8b86034b-c223-4c29-89ef-b03c24e736df
@bind rhf_plot html"""
<input type="checkbox" id="rhf">
RHF

"""

# ╔═╡ 8f15233b-3724-437f-a78c-460f69a41da6
	@bind mp2_plot html"""
	<input type="checkbox" id="mp2">
	MP2

	"""

# ╔═╡ 0d4ae0b1-6404-4471-818e-6cee43161b64
	@bind cc_plot html"""
	<input type="checkbox" id="mp2">
	CCSD

	"""

# ╔═╡ 850513be-8a24-4d4e-9d79-43062efaa50f
begin
	pt = plot(ylim=[-1.15, -0.80],legend=:topleft)
	xlabel!("H-H distance (Å)")
	ylabel!("Energy (Hartrees)")
	rhf_plot ? scatter!(dist, Eref, label="RHF", color=:steelblue2) : nothing
	mp2_plot ? scatter!(dist, Emp2, label="MP2",color=:seagreen) : nothing
	cc_plot ? scatter!(dist, Ecc, label="CCSD",color=:crimson) : nothing
	pt
end

# ╔═╡ Cell order:
# ╟─1ad27f33-fe6c-498a-b4e8-89948dd1f94b
# ╟─c51d70f8-9065-41f9-aea3-73f7c2aba7e7
# ╟─dc46fbbd-b96e-4fdd-9300-fa5ca87e86de
# ╟─97282505-47ee-4011-8198-bd4a7f8f75bf
# ╠═416d9c62-d2ef-11eb-3a40-450aee175208
# ╟─71915091-164d-440f-931e-7abaa1764e77
# ╠═2f8ca46a-3a92-4d0f-a28b-7e47fc577c80
# ╟─51444ecc-7031-4f9b-bb67-733bb6def9dc
# ╟─3a9d4c13-9d77-4deb-add2-b54e29d59934
# ╟─20b7f076-7177-41fe-b0d6-abd1b9211360
# ╠═913f76eb-3087-4a85-ae69-5dd7565964e9
# ╠═57f7f7c7-8a84-4fd0-b3f9-f49d375bf37c
# ╠═7388eaa1-011c-43df-9986-2f59ba259873
# ╠═77186eea-5b85-40c9-91bf-f48707752bb7
# ╠═ee4dbc82-6258-413c-aca5-1bd1e6ee02fc
# ╟─f17bc090-cad6-427a-ace9-7ed921f8368d
# ╠═aec67321-b943-478b-9765-8cb95c0f9d06
# ╠═15cfa11b-78e1-456c-82f2-02bcd1f8c6c8
# ╠═3685dd55-cb6a-4e0d-8a64-a76fd29d8e49
# ╟─b6e3d8a9-a125-417c-8af1-145b9050d406
# ╠═f41f12be-3755-4f73-a183-497ced9afa24
# ╠═6fb41a7c-2ae6-4b98-ba5a-fe122512364a
# ╠═90087552-51d4-42ef-8c3e-e5a61eec4fe1
# ╟─683b958d-7959-484e-baf1-9803094dd994
# ╠═42b34173-623d-4361-a2ba-b55b0cdfd69f
# ╠═165a0055-0195-42d8-8dd2-a3d07fe67f61
# ╟─9ab18d9f-28a2-43ae-a9df-012223dea282
# ╟─8b86034b-c223-4c29-89ef-b03c24e736df
# ╟─8f15233b-3724-437f-a78c-460f69a41da6
# ╟─0d4ae0b1-6404-4471-818e-6cee43161b64
# ╟─850513be-8a24-4d4e-9d79-43062efaa50f


text:
using Clang
using Documenter

makedocs(;
    modules=[Clang],
    repo="https://github.com/JuliaInterop/Clang.jl/blob/{commit}{path}#L{line}",
    sitename="Clang.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://JuliaInterop.github.io/Clang.jl",
        assets=String[],
    ),
    pages = [
        "Introduction" => "index.md",
        "Tutorial" => "tutorial.md",
        "API Reference" => "api.md",
    ],
)

deploydocs(;
    repo="github.com/JuliaInterop/Clang.jl.git",
)


text:
# module Bukdu

function Base.getproperty(c::C, prop::Symbol) where {C <: ApplicationController}
    if prop in (:params, :query_params, :path_params, :body_params)
        getfield(c.conn, prop)
    else
        getfield(c, prop)
    end
end

"""
    redirect_to(conn::Conn, path::String)
"""
function redirect_to(conn::Conn, path::String)
    body = """Redirecting to $path."""
    conn.request.response.status = 302 # 302 Found
    push!(conn.request.response.headers, Pair("Location", path))
    render(HTML, body)
end

# module Bukdu


text:
using JLD2, JUDI

JLD2.@load "256cube_v.jld2" v
idx_ucfmt = find_water_bottom((v.-3500f0).*(v.>3500f0))

K = zeros(Float32, size(v))

for i = 1:256
    println("i=",i)
    for j = 1:256
        K[:,:,1:idx_ucfmt[i,j]-1] = 1.03*1f-3*v[:,:,1:idx_ucfmt[i,j]-1].+15f0
        K[:,:,idx_ucfmt[i,j]:idx_ucfmt[i,j]+8] .= 1f-4
        K[:,:,idx_ucfmt[i,j]+9:end] = 1.03*1f-3*v[:,:,idx_ucfmt[i,j]+9:end].+200f0
    end
end

JLD2.@save "256K.jld2" K

max_stars_repo_name: slimgroup/TwoPhaseSpeedTest

max_issues_repo_name: slimgroup/TwoPhaseSpeedTest

max_forks_repo_name: slimgroup/TwoPhaseSpeedTest


text:
@testset "MarkdownDemoCard" begin
    # default behavior
    simplest = democard("simplest.md")
    @test simplest.cover === nothing
    @test simplest.id == "Simplest"
    @test simplest.path == "simplest.md"
    @test simplest.title == "Simplest"
    @test simplest.description == "This is the content"

    @testset "parse" begin
        @testset "title, id and description" begin
            # MarkdownDemoCard doesn't parse title from the markdown contents
            title_1 = MarkdownDemoCard("title_1.md")
            @test title_1.id == ".-Custom-Title"
            @test title_1.title == "1. Custom Title"
            @test title_1.description == "This is the content"
            @test title_1.hidden == false
            @test title_1.author == ""
            @test title_1.date == DateTime(0)

            title_2 = MarkdownDemoCard("title_2.md")
            @test title_2.id == "custom_id_2"
            @test title_2.title == "Custom Title 2"
            @test title_2.description == "This is the content"

            title_3 = MarkdownDemoCard("title_3.md")
            @test title_3.id == "Custom-Title-3-1"
            @test title_3.title == "Custom Title 3-1"
            @test title_3.description == "This is the content"

            title_4 = MarkdownDemoCard("title_4.md")
            @test title_4.id == "Custom-Title"
            @test title_4.title == "Custom Title"
            @test title_4.description == "This is the content"

            title_5 = MarkdownDemoCard("title_5.md")
            @test title_5.id == "Custom-Title"
            @test title_5.title == "Custom Title"
            @test title_5.description == "Custom Description"

            title_6 = MarkdownDemoCard("title_6.md")
            @test title_6.id == "custom_id"
            @test title_6.title == "Custom Title"
            @test title_6.description == "This is the content"
            @test title_6.hidden == false
            @test title_6.author == "Jane Doe; John Roe"
            @test title_6.date == DateTime("2020-01-31")
            @test title_6.cover == "https://juliaimages.org/latest/assets/logo.png"

            title_7 = MarkdownDemoCard("title_7.md")
            @test title_7.id == "custom_id"
            @test title_7.title == "Custom Title"
            @test title_7.description == "This is the content"

            description_1 = MarkdownDemoCard("description_1.md")
            @test description_1.id == "Custom-Title"
            @test description_1.title == "Custom Title"
            @test description_1.description == "this is a single line destiption that spans over multiple lines\n"

            description_2 = MarkdownDemoCard("description_2.md")
            @test description_2.id == "Custom-Title"
            @test description_2.title == "Custom Title"
            @test description_2.description == "this is a multi line\ndestiption that spans\nover multiple lines\n"
        end

        @testset "cover" begin
            cover_1 = MarkdownDemoCard("cover_1.md")
            @test cover_1.cover == nothing

            cover_2 = MarkdownDemoCard("cover_2.md")
            @test cover_2.cover == joinpath("..", "logo.png")

            cover_3 = MarkdownDemoCard("cover_3.md")
            @test cover_3.cover == "nonexistence.jpg"

            cover_4 = MarkdownDemoCard("cover_4.md")
            @test cover_4.cover == joinpath("..", "logo.png")
        end

        @testset "generate" begin
            page_dir = @suppress_err preview_demos("title_6.md", theme="grid", require_html=false)
            card_path = joinpath(page_dir, "markdown", "title_6.md")
            @test_reference joinpath(test_root, "references", "cards", "markdown.md") read(card_path, String) by=ignore_CR
        end
    end
end


text:
# const priλ = "Dir"
# const priλ = "SDM"
# const priλ = "SBM"
#
# const priQ = "Dir"
# const priQ = "SBM"
#
# const srnd = 22219
# const TT = 60
# const L = 10
#
# whichdat = "sim1_sets"
# whichdat = "sim2_sets"
# whichdat = "pinkSalmon_first30_K4"

### if ARGS used

const priλ = ARGS[1]
const priQ = ARGS[2]
const srnd = parse(Int64, ARGS[3])
const TT = parse(Int64, ARGS[4])
const L = parse(Int64, ARGS[5])
const whichdat = ARGS[6]

### after user selections:

using BSON
using StatsBase

# pkg> add "https://github.com/mheiner/SparseProbVec.jl"
using SparseProbVec

# pkg> add "https://github.com/mheiner/MTD.jl"
using MTD

using Random
using Dates

# srnd = 22219
Random.seed!(srnd)

### Read in data
simdata = occursin(r"sim\d+_sets", whichdat) # run validation at end?

if simdata
    BSON.@load "data/$(whichdat).bson" y X P
    const S = deepcopy(y[1])
else
    BSON.@load "data/$(whichdat).bson" S
end

### Setup
const SS = S[1:TT]
println(sort(unique(SS[(L+1):TT])))
println(counts(SS))
const K = maximum(S)

const n_keep = Int(2e3)
const n_burn = Int(200e3)
const thin = 200

const mesg = "$(whichdat)_MTD_T$(TT)_L$(L)_K$(K)_$(priλ)lam_$(priQ)Qmarg_$(srnd)"

### Prior

fieldnames(PriorMTD)
symm_prior = symmetricDirPrior_mtd(1.0, 1.0, L, K)

# β_λ = log(float(TT))
β_λ = sqrt(float(TT))
# β_λ = float(TT) / 4.0
# β_λ = 1.001

η_λ = 1.0e3
π_small_λ = 0.50
π_large_λ = 0.10
γ_λ, δ_λ = shape_Dir2GenDir(symm_prior[1])
δc_λ = del_correctionSBM(π_small_λ, π_large_λ, L)

# β_Q = float(TT) / 4.0 # for SDM on Q

η_Q = 1.0e3
π_small_Q = 0.75
π_large_Q = 0.10
γ_Q, δ_Q = shape_Dir2GenDir(symm_prior[2][:,1])
δc_Q = del_correctionSBM(π_small_Q, π_large_Q, K)

if priλ == "Dir"
        prior_λ = deepcopy(symm_prior[1])
    elseif priλ == "SDM"
        prior_λ = SparseDirMix(symm_prior[1], β_λ)
    elseif priλ == "SBM"
        prior_λ = SBMprior(L, η_λ, π_small_λ, π_large_λ, γ_λ, δc_λ.*δ_λ)
end

if priQ == "Dir"
        prior_Q = deepcopy(symm_prior[2])
    elseif priQ == "SDM"
        prior_Q = [ SparseDirMix(symm_prior[2][:,kk], β_Q) for kk = 1:K ]
    elseif priQ == "SBM"
        prior_Q = [ SBMprior(K, η_Q, π_small_Q, π_large_Q, γ_Q, δc_Q*δ_Q) for kk = 1:K ]
end

### Inits
fieldnames(ParamsMTD)
inits = ParamsMTD(
    SparseProbVec.rDirichlet(symm_prior[1], logout=true), # λ
    [ StatsBase.sample(1:L) for i = 1:(TT-L) ], # ζ
    reshape(vcat([ SparseProbVec.rDirichlet(ones(K), logout=true) for k in 1:K ]...), fill(K, 2)... )
  # log( transTens_MLE( count_trans_L(SS, K, 1) + 1.0 ) ) # lQ (mle)
  )

### Create model
fieldnames(ModMTD)
model = ModMTD(L, K, TT, SS,
  PriorMTD(prior_λ, prior_Q),
  inits)

logfilename = "postsim_progress/out_prog_$(mesg).txt"
timestart = Dates.now()

timemod!(5, model, 100, logfilename)

## Burn in
mcmc!(model, Int(0.15*n_burn)-500, save=false, report_filename=logfilename,
      thin=1, jmpstart_iter=5, report_freq=1000)
mcmc!(model, Int(0.85*n_burn), save=false, report_filename=logfilename,
      thin=1, jmpstart_iter=10, report_freq=10000)

## Estimate time remaining
etr(timestart, n_keep, thin, logfilename)

### MCMC
sims = mcmc!(model, n_keep, save=true, report_filename=logfilename,
             thin=thin, jmpstart_iter=10)

### Output
bson("./postsim/mcmc_$(mesg).bson", sims=deepcopy(sims),
    model=deepcopy(model), whichdat=deepcopy(whichdat),
    mesg=deepcopy(mesg))

### Send simulations to R
using RCall

nsim = length(sims)

R"rm(list=ls())"
priorinfo = "$(model.prior)"
sims_llik = [ sims[i][:llik] for i=1:nsim ]
sims_lam = permutedims(hcat([ exp.(sims[i][:lλ]) for i=1:nsim ]...))
sims_Q = permutedims(hcat([ exp.(vec(sims[i][:lQ])) for i=1:nsim ]...))

@rput sims_lam sims_Q sims_llik priorinfo K L TT mesg

R"ls()"
R"save.image(file=paste0('postsim/mcmc_', $(mesg), '.rda'))"

### GoF
if simdata

    using Random
    using Dates
    using Statistics

    modeltype = "MTD"
    report_filename = "forecast_validation/gof_$(mesg).txt"
    test_indx = 2

    report_file = open(report_filename, "a+")

    nprime = length(y[test_indx])
    niter = length(sims)
    nsim = 2000

    Random.seed!(srnd)
    simind = sort(StatsBase.sample(1:niter, nsim, replace=false))

    (FL1, MC, NLL, SQE, PL1, FOREC) = meanForecLoss(y[test_indx], X[test_indx][:,1:model.L],
        P[test_indx], nprime, lossL1, sims,
        model.TT, model.L, K,
        simind, modeltype=modeltype)

    PMforec = reshape(mean(FOREC, dims=1), (nprime, K) )
    PML1P = mean(abs.(PMforec .- P[test_indx])) # lower because of Jensen's inequality

    write(report_file, "L1 loss: $(round(100.0*mean(FL1), digits=2))\n" )
    write(report_file, "Misclass rate: $(round(100.0*mean(MC), digits=2))\n" )
    write(report_file, "NLL: $(round(mean(sum(NLL, dims=2)), digits=2))\n" )
    write(report_file, "SqErr: $(round(100.0*mean(SQE), digits=2))\n" )
    write(report_file, "postmean L1 loss, Ptrue: $(round(100.0*mean(PL1), digits=2))\n" )
    write(report_file, "L1 loss, Ptrue postmean forec: $(round(100.0*PML1P, digits=2))\n" )

    close(report_file)
end


text:
abstract type StructParent{T,N} <: AbstractArray{T,N} end
struct Struct{T} <: StructParent{T,1}
    x::Vector{T}
end
primitive type MyInt8 <: Integer 8 end
MyInt8(x::Integer) = Base.bitcast(MyInt8, convert(Int8, x))

myint = MyInt8(2)

const TypeAlias = Float32

# Methods and signatures
f1(x::Int) = 1
f1(x) = 2
f1(x::TypeAlias) = 3
# where signatures
f2(x::T) where T = -1
f2(x::T) where T<:Integer = T
f2(x::T) where Unsigned<:T<:Real = 0
f2(x::V) where V<:SubArray{T} where T = 2
f2(x::V) where V<:Array{T,N} where {T,N} = 3
f2(x::V) where V<:Base.ReshapedArray{T,N} where T where N = 4
# Varargs
f3(x::Int, y...) = 1
f3(x::Int, y::Symbol...) = 2
f3(x::T, y::U...) where {T<:Integer,U} = U
f3(x::Array{Float64,K}, y::Vararg{Symbol,K}) where K = K
# Default args
f4(x, y=0) = 1
f4(x, y::Int=0) = 2
f4(x::UInt, y="hello", z::Int=0) = 3
f4(x::Array{Float64,K}, y::Int=0) where K = K
# Keyword args
f5(x::Int8; y=0) = y
f5(x::Int16; y::Int=0) = 2
f5(x::Int32; y="hello", z::Int=0) = 3
f5(x::Int64;) = 4
f5(x::Array{Float64,K}; y::Int=0) where K = K
# Default and keyword args
f6(x, y="hello"; z::Int=0) = 1
# Destructured args
f7(x, (count, name)) = 1
# Return-type annotations
f8(x)::Int = 1
# generated functions
@generated function f9(x)
    if x <: Integer
        return :(x ^ 2)
    else
        return :(x)
    end
end
# Call overloading
(i::Struct)(::String) = i.x
(::Type{Struct{T}})(::Dict) where T = sizeof(T)

const first_two_funcs = (f1, f2)

# Conditional methods
if false
    ffalse(x) = 2
end
if true
    ftrue(x) = 3
end

if 0.8 > 0.2
    fctrue(x) = 1
else
    fcfalse(x) = 1
end

module Consts
export b1
b1 = true
b2 = false
g() = 2
end
using .Consts
if b1
    fb1true(x) = 1
else
    fb1false(x) = 1
end
if Consts.b2
    fb2true(x) = 1
else
    fb2false(x) = 1
end

if @isdefined(sum)
    fstrue(x) = 1
end

# Inner methods
function fouter(x)
    finner(::Float16) = 2x
    return finner(Float16(1))
end

## Evaled methods
for T in (Float32, Float64)
    @eval feval1(::$T) = 1
end
for T1 in (Float32, Float64), T2 in (Int8,)
    @eval feval2(::$T1, ::$T2) = 2
end
for f in (:length, :size)
    @eval Base.$f(i::Struct, args...) = nothing
end
for (T, v) in Dict(Float32=>4, Float64=>8)
    @eval nbytes(::Type{$T}) = $v
end
for x in (1, 1.1)
    @eval typestring(::$(typeof(x))) = $(string(typeof(x)))
end
for name in (:feval3,)
    _f = Symbol("_", name)
    @eval ($_f)(arg) = 3
end
const opnames = Dict{Symbol, Symbol}(:+ => :add, :- => :sub)
for op in [:+, :-, :max, :min]
    opname = get(opnames, op, op)
    @eval $(Symbol("feval_", opname, "!"))(var) = 1
end

# Methods with @isdefined
struct NoParam end
myeltype(::Type{Vector{T}}) where T = @isdefined(T) ? T : NoParam
paramtype(::Type{V}) where V<:Vector = isa(V, UnionAll) ? myeltype(Base.unwrap_unionall(V)) : myeltype(V)

## Submodules
module Inner
g() = 5
module InnerInner
g() = 6
end
end

module DatesMod
    abstract type Period end
end

struct Beat <: DatesMod.Period
    value::Int64
end

module Empty end


text:
module Materials
    __precompile__()
    #
    using LinearAlgebra
    #
    using StaticArrays
    using Parameters

    include("material_data_types.jl")
    include("material_lib.jl")
    include("eval_engineering_constants.jl")

    # const Material = Union{IsotropicMaterial, OrthotropicMaterial, TransverselyIsotropicMaterial}

    export AbstractMaterial, Material
    export IsotropicMaterial, OrthotropicMaterial, TransverselyIsotropicMaterial

end


text:
using OrdinaryDiffEq: WOperator, set_gamma!, calc_W, calc_W!
using OrdinaryDiffEq, LinearAlgebra, SparseArrays, Random, Test, DiffEqOperators

@testset "WOperator" begin
  Random.seed!(123)
  y = zeros(2); b = rand(2)
  mm = rand(2, 2)
  for _J in [rand(2, 2), Diagonal(rand(2))]
    _Ws = [-mm + 2.0 * _J, -mm/2.0 + _J]
    for inplace in (true, false), (_W, W_transform) in zip(_Ws, [false, true])
      W = WOperator(mm, 1.0, DiffEqArrayOperator(_J), inplace, transform=W_transform)
      set_gamma!(W, 2.0)
      @test convert(AbstractMatrix,W) ≈ _W
      @test W * b ≈ _W * b
      mul!(y, W, b); @test y ≈ _W * b
    end
  end
end

@testset "calc_W and calc_W!" begin
  A = [-1.0 0.0; 0.0 -0.5]; mm = [2.0 0.0; 0.0 1.0]
  u0 = [1.0, 1.0]; tmp = zeros(2)
  tspan = (0.0,1.0); dt = 0.01; dtgamma = 0.5dt
  concrete_W = -mm + dtgamma * A

  # Out-of-place
  fun = ODEFunction((u,p,t) -> A*u;
                    mass_matrix=mm,
                    jac=(u,p,t) -> A)
  integrator = init(ODEProblem(fun,u0,tspan), ImplicitEuler(); adaptive=false, dt=dt)
  W = calc_W(integrator, integrator.cache.nlsolver.cache, dtgamma, false)
  @test convert(AbstractMatrix, W) == concrete_W
  @test W \ u0 ≈ concrete_W \ u0

  # In-place
  fun = ODEFunction((du,u,p,t) -> mul!(du,A,u);
                    mass_matrix=mm,
                    jac_prototype=DiffEqArrayOperator(A))
  integrator = init(ODEProblem(fun,u0,tspan), ImplicitEuler(); adaptive=false, dt=dt)
  calc_W!(integrator.cache.nlsolver.cache.W, integrator, integrator.cache.nlsolver, integrator.cache, dtgamma, false)
  @test convert(AbstractMatrix, integrator.cache.nlsolver.cache.W) == concrete_W
  ldiv!(tmp, lu!(integrator.cache.nlsolver.cache.W), u0); @test tmp == concrete_W \ u0
end

@testset "Implicit solver with lazy W" begin
  A = sparse([-1.0 0.0; 0.0 -0.5])
  mm = sparse([2.0 0.0; 0.0 1.0])
  u0 = [1.0, 1.0]; tspan = (0.0,1.0)

  _f = (u,p,t) -> t*(A*u); _f_ip = (du,u,p,t) -> lmul!(t,mul!(du,A,u))
  fun1 = ODEFunction(_f; mass_matrix=mm)
  fun2 = ODEFunction(_f; mass_matrix=mm, jac=(u,p,t) -> t*A)
  fun1_ip = ODEFunction(_f_ip; mass_matrix=mm)
  fun2_ip = ODEFunction(_f_ip; mass_matrix=mm,
  jac_prototype=DiffEqArrayOperator(similar(A); update_func=(J,u,p,t) -> (J .= t .* A; J)))

  for Alg in [ImplicitEuler, Rosenbrock23, Rodas5]
    println(Alg)
    sol1 = solve(ODEProblem(fun1,u0,tspan), Alg(); adaptive=false, dt=0.01)
    sol2 = solve(ODEProblem(fun2,u0,tspan), Alg(); adaptive=false, dt=0.01)
    @test sol1(1.0) ≈ sol2(1.0)

    sol1_ip = solve(ODEProblem(fun1_ip,u0,tspan), Alg(); adaptive=false, dt=0.01)
    sol2_ip = solve(ODEProblem(fun2_ip,u0,tspan), Alg(linsolve=LinSolveFactorize(lu)); adaptive=false, dt=0.01)
    @test sol1_ip(1.0) ≈ sol2_ip(1.0) atol=1e-5
  end
end


text:
using IterTools
#My implementation of lexicographic less function, not sure if it is the most performant but it works
function lex_less(x, y)
    if (length(x) != length(y))
        return length(y) > length(x)
    else
        x_s = sort(x)
        y_s = sort(y)
        for i = 1:length(x)
            if (x_s[i] != y_s[i])
                return y_s[i] > x_s[i]
            end
        end
    end
end

#This function solves the vectorial problem:
#Multi dimensional models have no clue of the name of the underlying process and no clue of the name of themselves,
#A pricer implemented for a multidimensional model with a single name option makes little sense since there his no
#way to map the option to the right underlying.
#In order to solve this problem the pricer is implemented in a strange way:
#A vector of integer is computed in the following way:
#Assuming 3 underlyings
#1st element ---> option referred to model[1]
#2nd element ---> option referred to model[2]
#3rd element ---> option referred to model[3]
#4th element ---> option referred to (model[1],model[2])
#5th element ---> option referred to (model[1],model[3])
#6th element ---> option referred to (model[2],model[3])
#7th element ---> option referred to (model[1],model[2],model[3])
function compute_indices(x::Integer)
    X = collect(1:x)
    #Subsets will return the empty set as well
    f(y) = filter(x -> !(length(x) == 0), y)
    k = subsets(X) |> collect |> f
    out = sort(k, lt = (x, y) -> lex_less(x, y))

    return out
end

#Following function returns the names (duplicated as well) of the underlying names that must simulate.
#An underlying to be simulated must have an option in the portfolio that is referred to it (directly or in a basket).
function complete_2(payoffs_underlyings, models_underlyings)
    out = String[]
    for x_ in payoffs_underlyings
        for y_ in models_underlyings
            if (any(z -> z == x_, split(y_, "_")) || y_ == x_)
                push!(out, y_)
            end
        end
    end
    return out
end


text:
module RedirectRequest

import ..Layer, ..request
using ..URIs
using ..Messages
using ..Pairs: setkv
import ..Header
import ..@debug, ..DEBUG_LEVEL

"""
    request(RedirectLayer, method, ::URI, headers, body) -> HTTP.Response

Redirects the request in the case of 3xx response status.
"""
abstract type RedirectLayer{Next <: Layer} <: Layer{Next} end
export RedirectLayer

function request(::Type{RedirectLayer{Next}},
                 method::String, url::URI, headers, body;
                 redirect_limit=3, forwardheaders=true, kw...) where Next
    count = 0
    while true

        res = request(Next, method, url, headers, body; reached_redirect_limit=(count == redirect_limit), kw...)

        if (count == redirect_limit
        ||  !isredirect(res)
        ||  (location = header(res, "Location")) == "")
            return res
        end

        kw = merge(merge(NamedTuple(), kw), (parent = res,))
        oldurl = url
        url = absuri(location, url)
        if forwardheaders
            headers = filter(headers) do h
                # false return values are filtered out
                header, value = h
                if header == "Host"
                    return false
                elseif (header in SENSITIVE_HEADERS
                    && !isdomainorsubdomain(url.host, oldurl.host))
                    return false
                else
                    return true
                end
            end
        else
            headers = Header[]
        end

        @debug 1 "➡  Redirect: $url"

        count += 1
    end

    @assert false "Unreachable!"
end

const SENSITIVE_HEADERS = Set([
    "Authorization",
    "Www-Authenticate",
    "Cookie",
    "Cookie2"
])

function isdomainorsubdomain(sub, parent)
    sub == parent && return true
    endswith(sub, parent) || return false
    return sub[length(sub)-length(parent)] == '.'
end

end # module RedirectRequest


text:
export ZeDevice, properties, compute_properties, module_properties, memory_properties, memory_access_properties, cache_properties, image_properties, p2p_properties

struct ZeDevice
    handle::ze_device_handle_t

    driver::ZeDriver

    # only accept handles, don't convert
    ZeDevice(handle::ze_device_handle_t, driver::ZeDriver) = new(handle, driver)
end

Base.unsafe_convert(::Type{ze_device_handle_t}, dev::ZeDevice) = dev.handle

function Base.show(io::IO, ::MIME"text/plain", dev::ZeDevice)
    props = properties(dev)
    print(io, "ZeDevice(")
    if props.type == ZE_DEVICE_TYPE_GPU
        print(io, "GPU")
    elseif props.type == ZE_DEVICE_TYPE_FPGA
        print(io, "FPGA")
    end
    print(io, ", vendor ")
    show(io, props.vendorId)
    print(io, ", device ")
    show(io, props.deviceId)
    if props.subdeviceId !== nothing
        print(io, ", sub-device ")
        show(io, props.subdeviceId)
    end
    print(io, "): $(props.name)")
end

## properties

function properties(dev::ZeDevice)
    props_ref = Ref(ze_device_properties_t())
    zeDeviceGetProperties(dev, props_ref)

    props = props_ref[]
    return (
        type=props.type,
        vendorId=UInt16(props.vendorId),
        deviceId=UInt16(props.deviceId),
        flags=props.flags,
        subdeviceId=(props.flags&ZE_DEVICE_PROPERTY_FLAG_SUBDEVICE == 0) ? nothing : props.subdeviceId,
        coreClockRate=Int(props.coreClockRate),
        maxMemAllocSize=Int(props.maxMemAllocSize),
        maxHardwareContexts=Int(props.maxHardwareContexts),
        maxCommandQueuePriority=Int(props.maxCommandQueuePriority),
        numThreadsPerEU=Int(props.numThreadsPerEU),
        physicalEUSimdWidth=Int(props.physicalEUSimdWidth),
        numEUsPerSubslice=Int(props.numEUsPerSubslice),
        numSubslicesPerSlice=Int(props.numSubslicesPerSlice),
        numSlices=Int(props.numSlices),
        timerResolution=Int(props.timerResolution),
        timestampValidBits=Int(props.timestampValidBits),
        kernelTimestampValidBits=Int(props.kernelTimestampValidBits),
        uuid=Base.UUID(reinterpret(UInt128, [props.uuid.id...])[1]),
        name=String(UInt8[props.name[1:findfirst(isequal(0), props.name)-1]...]),
    )
end

function compute_properties(dev::ZeDevice)
    props_ref = Ref(ze_device_compute_properties_t())
    zeDeviceGetComputeProperties(dev, props_ref)

    props = props_ref[]
    return (
        maxTotalGroupSize=Int(props.maxTotalGroupSize),
        maxGroupSizeX=Int(props.maxGroupSizeX),
        maxGroupSizeY=Int(props.maxGroupSizeY),
        maxGroupSizeZ=Int(props.maxGroupSizeZ),
        maxGroupCountX=Int(props.maxGroupCountX),
        maxGroupCountY=Int(props.maxGroupCountY),
        maxGroupCountZ=Int(props.maxGroupCountZ),
        maxSharedLocalMemory=Int(props.maxSharedLocalMemory),
        subGroupSizes=Int.(props.subGroupSizes[1:props.numSubGroupSizes]),
    )
end

function module_properties(dev::ZeDevice)
    props_ref = Ref(ze_device_module_properties_t())
    zeDeviceGetModuleProperties(dev, props_ref)

    props = props_ref[]
    return (
        spirvVersionSupported=props.spirvVersionSupported==0 ? nothing : unmake_version(props.spirvVersionSupported),
        flags=props.flags,
        fp16flags=props.fp16flags,
        fp32flags=props.fp32flags,
        fp64flags=props.fp64flags,
        maxArgumentsSize=Int(props.maxArgumentsSize),
        printfBufferSize=Int(props.printfBufferSize),
        nativeKernelSupported=Base.UUID(reinterpret(UInt128, [props.nativeKernelSupported.id...])[1]),
    )
end

function memory_properties(dev::ZeDevice)
    count_ref = Ref{UInt32}(0)
    zeDeviceGetMemoryProperties(dev, count_ref, C_NULL)

    all_props = fill(ze_device_memory_properties_t(), count_ref[])
    zeDeviceGetMemoryProperties(dev, count_ref, all_props)

    return [(maxClockRate=Int(props.maxClockRate),
             maxBusWidth=Int(props.maxBusWidth),
             totalSize=Int(props.totalSize),
            ) for props in all_props[1:count_ref[]]]
end

function memory_access_properties(dev::ZeDevice)
    props_ref = Ref(ze_device_memory_access_properties_t())
    zeDeviceGetMemoryAccessProperties(dev, props_ref)

    props = props_ref[]
    return (
        hostAllocCapabilities=Int(props.hostAllocCapabilities),
        deviceAllocCapabilities=Int(props.deviceAllocCapabilities),
        sharedSingleDeviceAllocCapabilities=Int(props.sharedSingleDeviceAllocCapabilities),
        sharedCrossDeviceAllocCapabilities=Int(props.sharedCrossDeviceAllocCapabilities),
        sharedSystemAllocCapabilities=Int(props.sharedSystemAllocCapabilities),
    )
end

function cache_properties(dev::ZeDevice)
    count_ref = Ref{UInt32}(0)
    zeDeviceGetCacheProperties(dev, count_ref, C_NULL)

    all_props = fill(ze_device_cache_properties_t(), count_ref[])
    zeDeviceGetCacheProperties(dev, count_ref, all_props)

    return [(flags=props.flags,
             cacheSize=Int(props.cacheSize),
            ) for props in all_props[1:count_ref[]]]
end

function image_properties(dev::ZeDevice)
    props_ref = Ref(ze_device_image_properties_t())
    zeDeviceGetImageProperties(dev, props_ref)

    props = props_ref[]
    return (
        maxImageDims1D=Int(props.maxImageDims1D),
        maxImageDims2D=Int(props.maxImageDims2D),
        maxImageDims3D=Int(props.maxImageDims3D),
        maxImageBufferSize=Int(props.maxImageBufferSize),
        maxImageArraySlices=Int(props.maxImageArraySlices),
        maxSamplers=Int(props.maxSamplers),
        maxReadImageArgs=Int(props.maxReadImageArgs),
        maxWriteImageArgs=Int(props.maxWriteImageArgs),
    )
end

function p2p_properties(dev1, dev2::ZeDevice)
    props_ref = Ref(ze_device_p2p_properties_t())
    zeDeviceGetP2PProperties(dev1, dev2, props_ref)

    props = props_ref[]
    return (
        flags=props.flags,
    )
end

## device iteration

export devices

struct ZeDevices
    handles::Vector{ze_device_handle_t}

    driver::ZeDriver

    function ZeDevices(drv::ZeDriver)
        count_ref = Ref{UInt32}(0)
        zeDeviceGet(drv, count_ref, C_NULL)

        handles = fill(ze_device_handle_t(), count_ref[])
        zeDeviceGet(drv, count_ref, handles)

        new(handles, drv)
    end
end

devices(drv) = ZeDevices(drv)

Base.eltype(::ZeDevices) = ZeDevice

function Base.iterate(iter::ZeDevices, i=1)
    i >= length(iter) + 1 ? nothing : (ZeDevice(iter.handles[i], iter.driver), i+1)
end

Base.length(iter::ZeDevices) = length(iter.handles)

Base.IteratorSize(::ZeDevices) = Base.HasLength()


text:
#
# Restricted Two-body Problem Propagators
#

"""
Dynamics for ideal two-body numerical integration.
Compatable with `DifferentialEquations` solvers!
"""
function R2BPTic!(∂u, u, p, t)
    ∂u[1:3] .= u[4:6]
    ∂u[4:6] .= -p.μ .* (u[1:3] ./ norm(u[1:3])^3)
    return nothing
end

"""
Given a `CartesianOrbit`, return a `DifferentialEquations.ODEProblem` instance
that describes the Restricted Two-body Problem.
"""
function SciMLBase.ODEProblem(orbit::CartesianOrbit, Δt::Real = ustrip(timeunit(orbit.state), period(orbit)))

    return ODEProblem(
        AstrodynamicalModels.R2BPVectorField,
        MVector{6}(orbit.state.rv),
        (orbit.state.t, Δt),
        MVector{1}(orbit.system.μ)
    )

end

"""
Given a `CartesianOrbit`, return a `DifferentialEquations.ODEProblem` instance
that describes the Restricted Two-body Problem.
"""
SciMLBase.ODEProblem(orbit::CartesianOrbit, Δt::Unitful.Time) = ODEProblem(orbit, ustrip(timeunit(orbit.state), Δt))

"""
Uses `DifferentialEquations` solvers to propagate `orbit` Δt into the future.
All keyword arguments are passed directly to `DifferentialEquations.solve`.
"""
function propagate(orbit::CartesianOrbit,
                   Δt = period(orbit);
                   kwargs...)

    # Set default kwargs
    defaults = (;  reltol=1e-14, abstol=1e-14)
    options = merge(defaults, kwargs)

    # Initial conditions
    problem = ODEProblem(orbit, Δt)

    # Integrate!
    sols = solve(problem; options...)

    # Return PropagationResult structure
    if sols.retcode != :Success
        @warn "DifferentialEquations solvers returned code $(string(sols.retcode))."
    end

    return [
        CartesianOrbit(
            sols.u[i][1:3] * lengthunit(orbit.state),
            sols.u[i][4:6] * velocityunit(orbit.state),
            orbit.system, sols.t[i] * timeunit(orbit.state))
        for i ∈ 1:length(sols.t)
    ]
end

propagate(::KeplerianOrbit, args...; kwargs...) = throw(
    ArgumentError(
        "Keplerian orbit propagation isn't supported at this time! Please convert your orbit to a `CartesianOrbit` for propagation."
    )
)


text:
using GlobalSensitivity, QuasiMonteCarlo, Test, OrdinaryDiffEq

function ishi_batch(X)
    A= 7
    B= 0.1
    @. sin(X[1,:]) + A*sin(X[2,:])^2+ B*X[3,:]^4 *sin(X[1,:])
end
function ishi(X)
    A= 7
    B= 0.1
    sin(X[1]) + A*sin(X[2])^2+ B*X[3]^4 *sin(X[1])
end

function linear_batch(X)
    A= 7
    B= 0.1
    @. A*X[1,:]+B*X[2,:]
end
function linear(X)
    A= 7
    B= 0.1
    A*X[1]+B*X[2]
end

n = 600000
lb = -ones(4)*π
ub = ones(4)*π
sampler = SobolSample()
A,B = QuasiMonteCarlo.generate_design_matrices(n,lb,ub,sampler)

res1 = gsa(ishi,Sobol(order=[0,1,2]),A,B)
res2 = gsa(ishi_batch,Sobol(),A,B,batch=true)

@test res1.S1 ≈ [0.3139335358797363, 0.44235918402206326, 0.0, 0.0] atol=1e-4
@test res2.S1 ≈ [0.3139335358797363, 0.44235918402206326, 0.0, 0.0] atol=1e-4

@test res1.ST ≈ [0.5576009081644232, 0.44237102330046346, 0.24366241588532553, 0.0] atol=1e-4
@test res2.ST ≈ [0.5576009081644232, 0.44237102330046346, 0.24366241588532553, 0.0] atol=1e-4

@test res1.S2 ≈ [0.0 1.2954849537879917e-6 0.24368597775338205 6.64142747164482e-6; 0.0 0.0 4.279200031668718e-5 1.2542212940962112e-5; 0.0 0.0 0.0 -7.998213172266514e-7; 0.0 0.0 0.0 0.0] atol=1e-4

res1 = gsa(ishi,Sobol(order=[0,1,2], nboot = 20),A,B)
@test res1.S1 ≈ [0.3139335358797363, 0.44235918402206326, 0.0, 0.0] atol=1e-4
@test res1.ST ≈ [0.5576009081644232, 0.44237102330046346, 0.24366241588532553, 0.0] atol=1e-4
@test res1.S2 ≈ [0.0 1.2954849537879917e-6 0.24368597775338205 6.64142747164482e-6; 0.0 0.0 4.279200031668718e-5 1.2542212940962112e-5; 0.0 0.0 0.0 -7.998213172266514e-7; 0.0 0.0 0.0 0.0] atol=1e-4
@test res1.S1_Conf_Int ≈ [0.00013100970128286063, 0.00014730548523359544, 7.398816006175431e-5, 0.0] atol=1e-4
@test res1.ST_Conf_Int ≈ [ 5.657364947147881e-5, 0.00015856915718858496, 0.00012283019177515212, 0.0] atol=1e-4
@test res1.S2_Conf_Int ≈ [0.0 0.00019922618025106458 0.00017554669020070315 0.00020712452452973623; 0.0 0.0 0.00010571366158995006 0.00010243796353601678; 0.0 0.0 0.0 8.791383305689058e-5; 0.0 0.0 0.0 0.0] atol=1e-4

res1 = gsa(linear,Sobol(),A,B)
res2 = gsa(linear_batch,Sobol(),A,B,batch=true)

@test res1.S1 ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4
@test res2.S1 ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4

@test res1.ST ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4
@test res2.ST ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4

@test_nowarn gsa(linear,Sobol(order=[0,1,2], nboot = 100),A,B)

#=
library(sensitivity)
ishigami.fun <- function(X) {
  A <- 7
  B <- 0.1
  sin(X[, 1]) + A * sin(X[, 2])^2 + B * X[, 3]^4 * sin(X[, 1])
}
n <- 600000
X1 <- data.frame(matrix(runif(4 * n,-pi,pi), nrow = n))
X2 <- data.frame(matrix(runif(4 * n,-pi,pi), nrow = n))
sobol2007(ishigami.fun, X1, X2)
sobolSalt(ishigami.fun, X1, X2, scheme="A")
sobolSalt(ishigami.fun, X1, X2, scheme="B")

library(sensitivity)
ishigami.fun <- function(X) {
  A <- 7
  B <- 0.1
  A * X[, 1] + B * X[, 2]
}
n <- 6000000
X1 <- data.frame(matrix(runif(4 * n,-pi,pi), nrow = n))
X2 <- data.frame(matrix(runif(4 * n,-pi,pi), nrow = n))
sobol2007(ishigami.fun, X1, X2)
sobolSalt(ishigami.fun, X1, X2, scheme="A")
sobolSalt(ishigami.fun, X1, X2, scheme="B")
=#

function ishi_linear(X)
    A= 7
    B= 0.1
    [sin(X[1]) + A*sin(X[2])^2+ B*X[3]^4 *sin(X[1]),A*X[1]+B*X[2]]
end

function ishi_linear_batch(X)
    A= 7
    B= 0.1
    X1 = @. sin(X[1,:]) + A*sin(X[2,:])^2+ B*X[3,:]^4 *sin(X[1,:])
    X2 = @. A*X[1,:]+B*X[2,:]
    vcat(X1',X2')
end

res1 = gsa(ishi_linear,Sobol(),A,B)
res2 = gsa(ishi_linear_batch,Sobol(),A,B,batch=true)

# Now both tests together

@test res1.S1[1,:] ≈ [0.3139335358797363, 0.44235918402206326, 0.0, 0.0] atol=1e-4
@test res2.S1[1,:] ≈ [0.3139335358797363, 0.44235918402206326, 0.0, 0.0] atol=1e-4

@test res1.ST[1,:] ≈ [0.5576009081644232, 0.44237102330046346, 0.24366241588532553, 0.0] atol=1e-4
@test res2.ST[1,:] ≈ [0.5576009081644232, 0.44237102330046346, 0.24366241588532553, 0.0] atol=1e-4

@test res1.S1[2,:] ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4
@test res2.S1[2,:] ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4

@test res1.ST[2,:] ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4
@test res2.ST[2,:] ≈ [0.9997953478183109, 0.0002040399766938839, 0.0, 0.0] atol=1e-4

function f(du,u,p,t)
  du[1] = p[1]*u[1] - p[2]*u[1]*u[2] #prey
  du[2] = -p[3]*u[2] + p[4]*u[1]*u[2] #predator
end

u0 = [1.0;1.0]
tspan = (0.0,10.0)
p = [1.5,1.0,3.0,1.0]
prob = ODEProblem(f,u0,tspan,p)
t = collect(range(0, stop=10, length=200))

f1 = let prob = prob, t = t
    function (p)
        prob1 = remake(prob;p=p)
        sol = solve(prob1,Tsit5();saveat=t)
        return sol
    end
end

m = gsa(f1,Sobol(),[[1,5],[1,5],[1,5],[1,5]],N=100)
@test m isa GlobalSensitivity.SobolResult
m = gsa(f1,Sobol(order=[0,1,2], nboot = 10),[[1,5],[1,5],[1,5],[1,5]],N=100)
@test m isa GlobalSensitivity.SobolResult


text:
using Adapt
using Statistics
using Oceananigans.Grids
using Oceananigans.Grids: interior_parent_indices

"""
    struct AveragedField{X, Y, Z, A, G, N, O} <: AbstractReducedField{X, Y, Z, A, G, N}

Type representing an average over a field-like object.
"""
struct AveragedField{X, Y, Z, S, A, G, N, O} <: AbstractReducedField{X, Y, Z, A, G, N}
       data :: A
       grid :: G
       dims :: NTuple{N, Int}
    operand :: O
     status :: S

    function AveragedField{X, Y, Z}(data::A, grid::G, dims, operand::O;
                                    recompute_safely=true) where {X, Y, Z, A, G, O}

        dims = validate_reduced_dims(dims)
        validate_reduced_locations(X, Y, Z, dims)
        validate_field_data(X, Y, Z, data, grid)

        status = recompute_safely ? nothing : FieldStatus(0.0)

        S = typeof(status)
        N = length(dims)

        return new{X, Y, Z, S, A, G, N, O}(data, grid, dims, operand, status)
    end

    function AveragedField{X, Y, Z}(data::A, grid::G, dims, operand::O, status::S) where {X, Y, Z, A, G, O, S}
        return new{X, Y, Z, S, A, G, length(dims), O}(data, grid, dims, operand, status)
    end
end

operand_averaging_indices(avg_loc, op_loc, topo, N, H) = Colon() # fallback when dimension is non-reduced
operand_averaging_indices(::Type{Nothing}, op_loc, topo, N, H) = interior_parent_indices(op_loc, topo, N, H)

"""
    AveragedField(operand::AbstractField; dims, data=nothing, recompute_safely=false)

Returns an AveragedField.
"""
function AveragedField(operand::AbstractField; dims, data=nothing, recompute_safely=true)

    arch = architecture(operand)
    loc = reduced_location(location(operand), dims=dims)
    grid = operand.grid

    if isnothing(data)
        data = new_data(arch, grid, loc)
        recompute_safely = false
    end

    return AveragedField{loc[1], loc[2], loc[3]}(data, grid, dims, operand,
                                                 recompute_safely=recompute_safely)
end

"""
    compute!(avg::AveragedField)

Compute the average of `avg.operand` and store the result in `avg.data`.
"""
function compute!(avg::AveragedField, time=nothing)
    compute_at!(avg.operand, time)

    # Omit halo regions from operand on averaged dimension
    operand_parent = parent(avg.operand)

    Xa, Ya, Za = location(avg)
    Xo, Yo, Zo = location(avg.operand)

    i = operand_averaging_indices(Xa, Xo, topology(avg.grid, 1), avg.grid.Nx, avg.grid.Hx)
    j = operand_averaging_indices(Ya, Yo, topology(avg.grid, 2), avg.grid.Ny, avg.grid.Hy)
    k = operand_averaging_indices(Za, Zo, topology(avg.grid, 3), avg.grid.Nz, avg.grid.Hz)

    operand_sans_halos = view(operand_parent, i, j, k)

    mean!(avg.data.parent, operand_sans_halos)

    return nothing
end

compute_at!(avg::AveragedField{X, Y, Z, <:FieldStatus}, time) where {X, Y, Z} =
    conditional_compute!(avg, time)

#####
##### Adapt
#####

Adapt.adapt_structure(to, averaged_field::AveragedField{X, Y, Z}) where {X, Y, Z} =
    AveragedField{X, Y, Z}(Adapt.adapt(to, averaged_field.data),
                           nothing, averaged_field.dims, nothing, nothing)


text:
using Test, ArgParse
using SignalingDimension

try print_test_results catch err
    include("../utilities.jl")
    using .ScriptUtilities: print_test_results
end

# parsing command line arguments
arg_table = ArgParseSettings()
@add_arg_table! arg_table begin
    "-Y"
        help = "max number of outputs to consider"
        arg_type = Int64
        default = 12
    "--dir"
        help = "directory to print results .txt to"
        arg_type = String
        default = "./script/facet_verifications/results/"
    "--stdout"
        help = "Print output to STDOUT. If not flagged, output is printed to `.txt` file."
        action = :store_true
end
args = parse_args(ARGS, arg_table)

# verification script: verifies all k-guessing facets with Y ∈ [6, Y_max], d > 2, and k > 2
function verify_k_guessing_facet(Y_max :: Int64)
    @testset "Testing all k-guessing facets with `k > 2`, `d > 2`, and `Y ≤ $Y_max`." begin
        @testset "Scanning over all `d ∈ [3, $Y - 3]`" for Y in 6:Y_max
            @testset "Scanning over all `ε ∈ [3, $Y - $d + 1]`" for d in 3:Y-3
                @testset "Verifying N =  $Y, d = $d, k = $Y - $d" begin
                    k = Y - d
                    vertices = aff_ind_k_guessing_vertices(Y, d, k)
                    facet = k_guessing_facet(Y, d, k)

                    @test verify_facet(d, facet, vertices)
                end
            end
        end
    end
end

# outputting result to file or stdout
print_test_results(
    verify_k_guessing_facet,
    params=[args["Y"]],
    stdout=args["stdout"],
    dir=args["dir"]
)


text:
# MIT license
# Copyright (c) Microsoft Corporation. All rights reserved.
# See LICENSE in the project root for full license information.

""" Base class for defining points on a lattice. Lattice points are defined by basis vectors eᵢ and lattice indices indexᵢ:

point = Σᵢ indexᵢ*eᵢ

Example: define a 2D basis.

```julia
julia> a = LatticeBasis([1,2],[3,4])
LatticeBasis{2, Int64}(SVector{2, Int64}[[1, 2], [3, 4]])
```

Array indexing is used to generate a lattice point: a[1,2] = 1*[1,2] + 2*[3,4]

```julia
julia> a[1,2]
2-element SVector{2, Int64} with indices SOneTo(2):
  7
 10
```

Subtypes supporting the AbstractBasis interface should implement these functions:

Returns the neighbors in ring n surrounding centerpoint, excluding centerpoint
```
neighbors(::Type{B},centerpoint::Tuple{T,T},neighborhoodsize::Int) where{T<:Real,B<:AbstractBasis}
```
Returns the lattice basis vectors that define the lattice
```
basis(a::S) where{S<:AbstractBasis}
```
Returns the vertices of the unit polygon that tiles the plane for the basis
```
tilevertices(a::S) where{S<:AbstractBasis}
```
"""
abstract type AbstractBasis{N,T <: Real} end
export AbstractBasis

function Base.getindex(A::B1, indices::Vararg{Int,N}) where {N,T,B1 <: AbstractBasis{N,T}}
    return basismatrix(A) * SVector{N,Int}(indices)
end

#     temp::SVector{N,T} = (basis(A)*SVector{N,Int}(indices))::SVector{N,T}
#     return temp
# end

Base.setindex!(A::AbstractBasis{N,T}, v, I::Vararg{Int,N}) where {T,N} = nothing # can't set lattice points. Might want to throw an exception instead.

"""computes the tile vertices at the offset lattice coordinate."""
tilevertices(latticecoords::Union{AbstractVector,NTuple},lattice::AbstractBasis) = lattice[latticecoords...] .+ tilevertices(lattice)

struct LatticeBasis{N,T <: Real} <: AbstractBasis{N,T}
    basisvectors::SMatrix{N,N,T}

    """ Convenience constructor that lets you use tuple arguments to describe the basis instead of SVector

    Example:
    ```
    basis = LatticeBasis{2}((1.0,0.0),(0.0,1.0)) #basis for rectangular lattice
    ```
    This allocates 48 bytes for a 2D basis in Julia 1.6.2. It  shouldn't allocate anything but not performance critical.
    """
    function LatticeBasis(vectors::Vararg{NTuple{N,T},N}) where {T,N}
        temp = MMatrix{N,N,T}(undef)

        for (j, val) in pairs(vectors)
            for i in 1:N
                temp[i,j] = val[i]
            end
        end

        return new{N,T}(SMatrix{N,N,T}(temp))
    end

     LatticeBasis(vectors::SMatrix{N,N,T}) where {N,T <: Real} = new{N,T}(vectors)
end
export LatticeBasis

function basismatrix(a::LatticeBasis{N,T})::SMatrix{N,N,T,N * N} where {N,T}
    return a.basisvectors
end

"""Can access any point in a lattice so the range of indices is unlimited"""
function Base.size(a::LatticeBasis{N,T}) where {N,T}
    return ntuple((i) -> Base.IsInfinite(), N)
end

# These functions are used to find the lattice tiles that lie inside a polygonal shape

"""computes the matrix that transforms the basis set into the canonical basic vectors, eᵢ"""
basisnormalization(a::Repeat.AbstractBasis) = Matrix(inv(Repeat.basismatrix(a))) # unfortunately LazySets doesn't work with StaticArrays. If you return a static Array is messes with their functions.

"""warp the matrix into a coordinate frame where the basis vectors are canonical, eᵢ"""
function normalizedshape(a::Repeat.AbstractBasis, shape::LazySets.VPolygon)
    normalizer = basisnormalization(a)
    return normalizer * shape
end
export normalizedshape

""" compute a conservative range of lattice indices that might intersect containingshape"""
function latticebox(containingshape::LazySets.VPolygon, lattice::Repeat.AbstractBasis)
    warpedshape = normalizedshape(lattice, containingshape)
    box = LazySets.interval_hull(warpedshape)
    return LazySets.Hyperrectangle(round.(box.center), (1, 1) .+ ceil.(box.radius)) # center the hyperrectangle on an integer point and enlarge the radius to take this into account
end
export latticebox

"""Returns the lattice coordinates of all tiles that lie at least partly inside containingshape.

Compute the lattice tiles with non-zero intersection with containingshape. First compute a transformation that maps the lattice basis vectors to canonical unit basis vectors eᵢ (this is the inverse of the lattice basis matrix). Then transform containingshape into this coordinate frame and compute a bounding box. Unit steps along the coordinate axes in this space represent unit *lattice* steps in the original space. This makes it simple to determine coordinate bounds in the original space. Then test for intersection in the original space."""
function tilesinside(containingshape::LazySets.VPolygon, lattice::Repeat.AbstractBasis{2,T}) where {T}
    coordtype = Int64

    box = latticebox(containingshape, lattice)

    coords = coordtype.(box.radius)
    tilevertices = LazySets.VPolygon(Matrix(Repeat.tilevertices(lattice))) # VPolygon will accept StaticArrays but other LazySets function will barf.
    result = Matrix{coordtype}(undef, 2, 0)

    for i in -coords[1]:coords[1]
        for j in -coords[2]:coords[2]
            offsetcoords =   (Int64.(box.center) + [i,j])
            center = Vector(lattice[offsetcoords...]) #[] indexer for lattices returns SVector which LazySets doesn't handle. Convert to conventional Vector.
            offsethex = LazySets.translate(tilevertices, center)

            if !isempty(offsethex ∩ containingshape)
                 result = hcat(result, offsetcoords)
            end
        end
    end
    return result
end
export tilesinside

"""The vertices of the containing shape are the columns of the matrix containingshape"""
tilesinside(containingshape::AbstractMatrix,lattice::Repeat.AbstractBasis) = tilesinside(LazySets.VPolygon(containingshape), lattice)

tilesinside(xmin::T,ymin::T,xmax::T,ymax::T,lattice::Repeat.AbstractBasis) where {T <: AbstractFloat} = tilesinside(SMatrix{2,4,T}(xmin, ymin, xmin, ymax, xmax, ymax, xmax, ymin), lattice)

using Plots

""" to see what the objects look like in the warped coordinate frame use inv(basismatrix(lattice)) as the transform"""
function plotall(containingshape, lattice, transform=[1.0 0.0;0.0 1.0])
    temp = Vector{SMatrix}(undef, 0)

    for coordinate in eachcol(tilesinside(containingshape, lattice))
        println(typeof(coordinate))
        push!(temp, tilevertices(coordinate, lattice))
    end

    tiles = LazySets.VPolygon.(temp)
    for tile in tiles
        plot!(transform * tile, aspectratio=1)
    end
    plot!(containingshape, aspectratio=1)
end
export plotall

function testtilesinside()
    # poly = LazySets.VPolygon([-3.0 3.0 3.0; -3.0 3.0 -3.0])
    hex = HexBasis1()
    poly = LazySets.VPolygon(3 * tilevertices(hex))
    tilecoords = tilesinside(poly, hex)
     insidetiles = LazySets.VPolygon.([tilevertices(x,hex) for x in eachcol(tilecoords)])
    # plotall(poly, hex)
    for tile in insidetiles
        plot!(tile)
end
end
export testtilesinside

"""This function returns the radius of the longest basis vector of the lattice cluster. Most lattices defined in this project have symmetric basis vectors so the radii of all basis vectors will be identical. This function is used in determing which clusters "fit" within the eye pupil."""
function euclideandiameter(basismatrix::SMatrix)
    maximum(norm.([basismatrix[:,i] for i in 1:size(basismatrix)[2]]))
end
export euclideandiameter

""" Lattic bases can have real basis vectors. This returns the maximum Euclidean distance between lattice basis center points."""
euclideandiameter(a::Repeat.AbstractBasis) =  euclideandiameter(Repeat.basismatrix(a))

"""Only will work properly if lattice basis matrix contains only integer or rational terms. Returns the integer lattice coords of point in the given basis if the point is in the span of latticebasis. Otherwise returns nothing"""
function latticepoint(lattice_basis_matrix::AbstractMatrix, origin, point)
    Ainv = inv(Rational.(lattice_basis_matrix))
    b = [(point .- origin)...]
    x = Ainv * b
    if reduce(&, (1, 1) .== denominator.(x))
        return Integer.(x)
    else
        return nothing
    end
end
export latticepoint

latticepoint(latticebasis::AbstractBasis,origin,point) = latticepoint(basismatrix(latticebasis),origin,point)


text:
"""
Transient analysis for CTMC
"""

"""
mexp(Q, x, t; transpose = :N, ufact = 1.01, eps = 1.0e-8, rmax = 500)

Compute the probability vector for CTMC.

exp(tr(Q)*t) * x

Parameters:
- Q: CTMC Kernel
- x: Array
- t: time
- transpose: forward or backward
- ufact: uniformization factor
- eps: tolerance error for Poisson p.m.f.
- rmax: The maximum number of uniformization steps

Return value:
- probability vector
"""

function mexp(Q::AbstractMatrix{Tv}, x::ArrayT, t::Tv;
    transpose::Symbol=:N, ufact::Tv=Tv(1.01), eps::Tv=Tv(1.0e-8), rmax=500) where {Tv, ArrayT <: AbstractArray{Tv}}
    m, n = size(Q)
    @assert m == n
    P, qv = unif(Q, ufact)
    right = rightbound(qv*t, eps)
    @assert right <= rmax "Time interval is too large. t or rmax should be changed: right = $right (rmax: $rmax)."
    weight, poi = poipmf(qv*t, right, left = 0)
    y = zero(x)
    unifstep!(transpose, P, poi, (0, right), weight, copy(x), y)
    return y
end

"""
mexpc(Q, x, t; transpose = :N, ufact = 1.01, eps = 1.0e-8, rmax = 500)

Compute the probability vector for CTMC and the cumulative value.

exp(tr(Q)*t) * x
int_0^t exp(tr(Q)*u) * x du

Parameters:
- Q: CTMC Kernel
- x: Array
- t: time
- transpose: forward or backward
- ufact: uniformization factor
- eps: tolerance error for Poisson p.m.f.
- rmax: The maximum number of uniformization steps

Return value (tuple)
- probability vector
- cumulative value
"""

function mexpc(Q::AbstractMatrix{Tv}, x::ArrayT, t::Tv;
    transpose::Symbol=:N, ufact::Tv=Tv(1.01), eps::Tv=Tv(1.0e-8), rmax=500) where {Tv, ArrayT <: AbstractArray{Tv}}
    m, n = size(Q)
    @assert m == n
    P, qv = unif(Q, ufact)
    right = rightbound(qv*t, eps) + 1
    @assert right <= rmax "Time interval is too large. t or rmax should be changed: right = $right (rmax: $rmax)."
    weight, poi, cpoi = cpoipmf(qv*t, right, left = 0)
    y = zero(x)
    cy = zero(x)
    cunifstep!(transpose, P, poi, cpoi, (0, right), weight, qv*weight, copy(x), y, cy)
    return y, cy
end

"""
mexp(Q, x, ts; transpose = :N, ufact = 1.01, eps = 1.0e-8, rmax = 500)

Compute the probability vector for CTMC for time series

exp(tr(Q)*t) * x for t = ts

Parameters:
- Q: CTMC Kernel
- x: Array
- ts: time series
- transpose: forward or backward
- ufact: uniformization factor
- eps: tolerance error for Poisson p.m.f.
- rmax: The maximum number of uniformization steps

Return value:
- probability vector
"""

function mexp(Q::AbstractMatrix{Tv}, x::ArrayT, ts::AbstractVector{Tv};
    transpose::Symbol=:N, ufact::Tv=Tv(1.01), eps::Tv=Tv(1.0e-8), rmax=500) where {Tv,ArrayT<:AbstractArray{Tv}}
    m, n = size(Q)
    @assert m == n
    dt, maxt = itime(sort(ts))
    P, qv = unif(Q, ufact)
    right = rightbound(qv*maxt, eps)
    @assert right <= rmax "Time interval is too large. t or rmax should be changed: right = $right (rmax: $rmax)."
    prob = Vector{Tv}(undef, right+1)
    # result = Vector{ArrayT}(undef, length(dt))
    result = Vector{Any}(undef, length(dt)) # TODO: memory usage?
    y0 = copy(x)
    for k = eachindex(dt)
        right = rightbound(qv*dt[k], eps)
        weight = poipmf!(qv*dt[k], prob; left=0, right=right)
        y1 = zero(y0)
        unifstep!(transpose, P, prob, (0, right), weight, y0, y1)
        result[k] = y1
        y0 .= y1
    end
    return result
end

"""
mexpc(Q, x, ts; transpose = :N, ufact = 1.01, eps = 1.0e-8, rmax = 500)

Compute the probability vector for CTMC and the cumulative value for time series.

exp(tr(Q)*t) * x for t = ts
int_0^t exp(tr(Q)*u) * x du for t = ts

Parameters:
- Q: CTMC Kernel
- x: Array
- ts: time series
- transpose: forward or backward
- ufact: uniformization factor
- eps: tolerance error for Poisson p.m.f.
- rmax: The maximum number of uniformization steps

Return value (tuple)
- probability vector
- cumulative value
"""

function mexpc(Q::AbstractMatrix{Tv}, x::ArrayT, ts::AbstractVector{Tv};
    transpose::Symbol=:N, ufact::Tv=Tv(1.01), eps::Tv=Tv(1.0e-8), rmax=500) where {Tv,ArrayT<:AbstractArray{Tv}}
    m, n = size(Q)
    @assert m == n
    dt, maxt = itime(ts)
    P, qv = unif(Q, ufact)
    right = rightbound(qv*maxt, eps) + 1
    @assert right <= rmax "Time interval is too large. t or rmax should be changed: right = $right (rmax: $rmax)."
    prob = Vector{Tv}(undef, right+1)
    cprob = Vector{Tv}(undef, right+1)
    # result = Vector{ArrayT}(undef, length(dt))
    # cresult = Vector{ArrayT}(undef, length(dt))
    result = Vector{Any}(undef, length(dt)) # TODO: memory usage?
    cresult = Vector{Any}(undef, length(dt)) # TODO: memory usage?
    y0 = copy(x)
    cy = zero(x)
    tmp = similar(x)
    for k = eachindex(dt)
        right = rightbound(qv*dt[k], eps) + 1
        weight = cpoipmf!(qv*dt[k], prob, cprob; left=0, right=right)
        y1 = zero(y0)
        tmp .= Tv(0)
        cunifstep!(transpose, P, prob, cprob, (0, right), weight, qv*weight, y0, y1, tmp)
        cy .+= tmp
        result[k] = copy(y1)
        cresult[k] = copy(cy)
        y0 .= y1
    end
    return result, cresult
end


text:
@doc raw"""
    rotate(p::Particles, rotmat::AbstractMatrix{<:Real})
    rotate!(p::Particles, rotmat::AbstractMatrix{<:Real})

Rotates the particles `p` by the ``3 \times 3`` rotation matrix `rotmat`.

The non-in-place version `rotate` creates a copy of the particles with only new pointers to the
rotated properties (generally "POS" and "VEL").
"""
function rotate(p::Particles, rotmat::AbstractMatrix{<:Real})
    @assert size(rotmat) == (3, 3)

    pc = copy(p)

    # only rotate existing quantities (currently only pos and vel)
    for prop in intersect(keys(pc), [:pos, :vel])
        pc[prop] = rotate(pc[prop], rotmat)
    end

    return pc
end

function LinearAlgebra.rotate!(p::Particles, rotmat::AbstractMatrix{<:Real})
    @assert size(rotmat) == (3, 3)

    # only rotate existing quantities (currently only pos and vel)
    for prop in intersect(keys(p), [:pos, :vel])
        rotate!(p[prop], rotmat)
    end

    return p
end

@doc raw"""
    rotate(g::AbstractGalaxy, rotmat::AbstractMatrix{<:Real})
    rotate!(g::AbstractGalaxy, rotmat::AbstractMatrix{<:Real})

Rotates the galaxy `g` by the ``3 \times 3`` rotation matrix `rotmat`.

The non-in-place version `rotate` creates a copy of the galaxy with only new pointers to the
rotated quantites (generally "POS" and "VEL") and copies of the particle Dicts.
"""
function rotate(g::AbstractGalaxy, rotmat::AbstractMatrix{<:Real})
    @assert size(rotmat) == (3, 3)

    gc = copy(g)

    # do for stars, dm, etc.
    for p in values(gc)
        # only rotate existing quantities (currently only pos and vel)
        for prop in intersect(keys(p), [:pos, :vel])
            p[prop] = rotate(p[prop], rotmat)
        end
    end

    return gc
end

function LinearAlgebra.rotate!(g::AbstractGalaxy, rotmat::AbstractMatrix{<:Real})
    @assert size(rotmat) == (3, 3)
    # do for stars, dm, etc.
    for p in values(g)
        rotate!(p, rotmat)
    end

    return g
end

"""
    rotate(vals::AbstractMatrix{<:Number}, rotmat::AbstractMatrix{<:Real})
    rotate!(vals::AbstractMatrix{<:Number}, rotmat::AbstractMatrix{<:Real})

Returns `vals` rotated by the rotation matrix `rotmat`. Works for any dimensions.
"""
rotate(vals::AbstractMatrix{<:Number}, rotmat::AbstractMatrix{<:Real}) = rotmat * vals

function LinearAlgebra.rotate!(vals::AbstractMatrix{<:Number}, rotmat::AbstractMatrix{<:Real})
    vals .= rotate(vals, rotmat)
    return vals
end

"""
    rotate_edgeon(p::Particles; axis_ratios::Bool=false, kwargs...)
    rotate_edgeon!(p::Particles; axis_ratios::Bool=false, kwargs...)

Rotates the particles edge-on using the given shape determination algorithm and `radius` (see [`rotation_matrix_edgeon`](@ref) for all keyword parameters).

Keyword parameters
- `axis_ratios::Bool`: `true` to return tuple `(p, q, s)` with particles and axis ratios
"""
function rotate_edgeon(p::Particles; axis_ratios::Bool=false, kwargs...)
    Q⁻¹, q, s = rotation_matrix_edgeon(p; kwargs...)

    pc = rotate(p, Q⁻¹)

    if axis_ratios
        return pc, q, s
    end

    return pc
end

function rotate_edgeon!(p::Particles; axis_ratios::Bool=false, kwargs...)
    Q⁻¹, q, s = rotation_matrix_edgeon(p; kwargs...)

    rotate!(p, Q⁻¹)

    if axis_ratios
        return p, q, s
    end

    return p
end

"""
    rotate_edgeon(g::AbstractGalaxy, ptype::Symbol=:stars; axis_ratios::Bool=false, kwargs...)
    rotate_edgeon!(g::AbstractGalaxy, ptype::Symbol=:stars; axis_ratios::Bool=false, kwargs...)

Rotates the galaxy's particles edge-on for the given particle type `ptype` (`:stars`, `:dm`, `:gas`, etc.)
using the given shape determination algorithm and `radius`
(see [`rotation_matrix_edgeon`](@ref) for all keyword parameters).

Keyword parameters
- `axis_ratios::Bool`: `true` to return tuple `(g, q, s)` with galaxy and axis ratios
"""
function rotate_edgeon(g::AbstractGalaxy, ptype::Symbol=:stars; axis_ratios::Bool=false, kwargs...)
    Q⁻¹, q, s = rotation_matrix_edgeon(g[ptype]; kwargs...)

    gc = rotate(g, Q⁻¹)

    if axis_ratios
        return gc, q, s
    end

    return gc
end

function rotate_edgeon!(g::AbstractGalaxy, ptype::Symbol=:stars; axis_ratios::Bool=false, kwargs...)
    Q⁻¹, q, s = rotation_matrix_edgeon(g[ptype]; kwargs...)

    rotate!(g, Q⁻¹)

    if axis_ratios
        return g, q, s
    end

    return g
end

"""
    translate(p::Particles, x⃗::AbstractVector{<:Number}, prop::Symbol=:pos)
    translate!(p::Particles, x⃗::AbstractVector{<:Number}, prop::Symbol=:pos)

Translates the particles `p` by the vector `x⃗`.

The non-in-place version `translate` creates a copy of the particles with only new pointers to the
property.
"""
function translate(p::Particles, x⃗::AbstractVector{<:Number}, prop::Symbol=:pos)
    @assert length(x⃗) == 3

    pc = copy(p)

    if haskey(pc, prop)
        pc[prop] = pc[prop] .+ x⃗
    end

    return pc
end

function translate!(p::Particles, x⃗::AbstractVector{<:Number}, prop::Symbol=:pos)
    @assert length(x⃗) == 3

    if haskey(p, prop)
        p[prop] .+= x⃗
    end

    return p
end

"""
    translate(g::AbstractGalaxy, x⃗::AbstractVector{<:Number}, prop::Symbol=:pos)
    translate!(g::AbstractGalaxy, x⃗::AbstractVector{<:Number}, prop::Symbol=:pos)

Translates the galaxy `g` by the vector `x⃗`.

The non-in-place version `translate` creates a copy of the galaxy with only new pointers to the
property and copies of the particle Dicts.
"""
function translate(g::AbstractGalaxy, x⃗::AbstractVector{<:Number}, prop::Symbol=:pos)
    @assert length(x⃗) == 3

    gc = copy(g)

    # do for stars, dm, etc.
    for p in values(gc)
        if haskey(p, prop)
            p[prop] = p[prop] .+ x⃗
        end
    end

    return gc
end

function translate!(g::AbstractGalaxy, x⃗::AbstractVector{<:Number}, prop::Symbol=:pos)
    @assert length(x⃗) == 3

    # do for stars, dm, etc.
    for p in values(g)
        translate!(p, x⃗, prop)
    end

    return g
end

@doc raw"""
    center_of_mass_iterative(
        pos::AbstractMatrix{<:Number},
        mass::AbstractVector{<:Number},
        r_start::Number;
        shrinking_factor::Real=0.025,
        limit_fraction::Real=0.01,
        limit_number::Integer=1000,
        limit_radius::Number=0,
    )

Returns the center of mass by applying the shrinking sphere method
([Power et al. 2003](https://arxiv.org/abs/astro-ph/0201544)).

The sphere starts centered at the coordinate system's origin with a radius of `r_start` and
shrinks by `shrinking_factor` every iteration. The iteration is stopped when the sphere
contains fewer than `limit_number` and fewer than `limit_fraction` times the number of particles in the
starting sphere, whichever is smaller. In addition, the iteration is stopped when the radius
of the sphere drops below `limit_radius`. Note that the unit of `limit_radius` should be the
same as that of `r_start`.

The radius of the sphere in iteration step ``i`` is ``r_0 (1 - f_\mathrm{shrink})^i``.
"""
function center_of_mass_iterative(
    pos::AbstractMatrix{T},
    mass::AbstractVector{<:Number},
    r_start::R;
    shrinking_factor::S=0.025,
    limit_fraction::Real=0.01,
    limit_number::Integer=1000,
    limit_radius::Number=zero(R),
) where {T<:Number,R<:Number,S<:Real}
    r₀ = zeros(T, 3)
    r²_max = r_start^2
    r²_limit = (limit_radius / (one(S) - shrinking_factor))^2

    # mask particles in initial sphere
    r² = r²_sphere(pos .- r₀)
    mask = r² .≤ r²_max
    n = count(mask)

    # get number limit for when the algorithm should stop
    nlimit = min(limit_number, ceil(Int, limit_fraction * n))
    while n ≥ nlimit && r²_max ≥ r²_limit
        r₀ .= @views dropdims(sum(pos[:, mask] .* mass[mask]'; dims=2); dims=2) ./ sum(mass[mask])
        r²_max *= (one(S) - shrinking_factor)^2

        r² .= r²_sphere(pos .- r₀)
        mask .= r² .≤ r²_max
        n = count(mask)
    end

    return r₀
end

"""
    center_of_mass_iterative(p::Particles, r_start::Number; kwargs...)

Returns the center of mass of the particles.
"""
function center_of_mass_iterative(p::Particles, r_start::Number; kwargs...)
    center_of_mass_iterative(p.pos, p.mass, r_start; kwargs...)
end

"""
    center_of_mass_iterative(g::AbstractGalaxy, r_start::Number, ptype::Symbol=:stars; kwargs...)

Returns the center of mass of the galaxy's particles of type `ptype` (`:stars`, `:dm`, `:gas`, etc.).
"""
function center_of_mass_iterative(g::AbstractGalaxy, r_start::Number, ptype::Symbol=:stars; kwargs...)
    center_of_mass_iterative(g[ptype], r_start; kwargs...)
end

"""
    translate_to_center_of_mass_iterative(p::Particles, r_start::Number; kwargs...)
    translate_to_center_of_mass_iterative!(p::Particles, r_start::Number; kwargs...)

Translates the particles to their center of mass of the particles (see [`center_of_mass_iterative`](@ref)).
"""
function translate_to_center_of_mass_iterative(p::Particles, r_start::Number; kwargs...)
    x⃗ = center_of_mass_iterative(p, r_start; kwargs...)

    return translate(p, -x⃗)
end

function translate_to_center_of_mass_iterative!(p::Particles, r_start::Number; kwargs...)
    x⃗ = center_of_mass_iterative(p, r_start; kwargs...)

    return translate!(p, -x⃗)
end

"""
    translate_to_center_of_mass_iterative(g::AbstractGalaxy, r_start::Number, ptype::Symbol=:stars; kwargs...)
    translate_to_center_of_mass_iterative!(g::AbstractGalaxy, r_start::Number, ptype::Symbol=:stars; kwargs...)

Translates the particles to their center of mass of the particles (see [`center_of_mass_iterative`](@ref)).
"""
function translate_to_center_of_mass_iterative(
    g::AbstractGalaxy,
    r_start::Number,
    ptype::Symbol=:stars;
    kwargs...,
)
    x⃗ = center_of_mass_iterative(g, r_start, ptype; kwargs...)

    return translate(g, -x⃗)
end

function translate_to_center_of_mass_iterative!(
    g::AbstractGalaxy,
    r_start::Number,
    ptype::Symbol=:stars;
    kwargs...,
)
    x⃗ = center_of_mass_iterative(g, r_start, ptype; kwargs...)

    return translate!(g, -x⃗)
end

"""
    center_of_velocity(
        pos::AbstractMatrix{<:Number},
        vel::AbstractMatrix{<:Number},
        mass::AbstractVector{<:Number},
        radius::Number;
        p::Real=0.9,
        r₀::Union{AbstractVector{<:Number},Nothing}=nothing,
    )

Returns mass-weighted mean velocity within `radius` around `r₀`, only considering absolute velocities
smaller than the `p`th quantile (any value between 0 and 1) when centered around the median velocity.
"""
function center_of_velocity(
    pos::AbstractMatrix{<:Number},
    vel::AbstractMatrix{<:Number},
    mass::AbstractVector{<:Number},
    radius::Number;
    p::Real=0.9,
    r₀::Union{AbstractVector{<:Number},Nothing}=nothing,
)
    if isnothing(r₀)
        r₀ = zeros(eltype(pos), 3)
    end

    r² = r²_sphere(pos .- r₀)
    mask = r² .≤ radius^2
    v² = @views r²_sphere(vel[:, mask] .- median(vel[:, mask]; dims=2))
    v²_max = quantile(v², p)
    mask2 = v² .≤ v²_max

    return @views dropdims(sum(vel[:, mask][:, mask2] .* mass[mask][mask2]'; dims=2); dims=2) ./
                  sum(mass[mask][mask2])
end

"""
    center_of_velocity(p::Particles, radius::Number; kwargs...)

Returns the center of velocity of the particles. Requires "POS", "VEL", and "MASS" to be available.
"""
function center_of_velocity(p::Particles, radius::Number; kwargs...)
    center_of_velocity(p.pos, p.vel, p.mass, radius; kwargs...)
end

"""
    center_of_velocity(g::AbstractGalaxy, radius::Number, ptype::Symbol=:stars; kwargs...)

Returns the center of velocity of the galaxy's particles of type `ptype` (`:stars`, `:dm`, `:gas`, etc.).
Requires "POS", "VEL", and "MASS" of the particle type to be available.
"""
function center_of_velocity(g::AbstractGalaxy, radius::Number, ptype::Symbol=:stars; kwargs...)
    center_of_velocity(g[ptype], radius; kwargs...)
end

"""
    translate_to_center_of_velocity(p::Particles, radius::Number; kwargs...)
    translate_to_center_of_velocity!(p::Particles, radius::Number; kwargs...)

Translates the particles to their center of velocity of the particles (see [`center_of_velocity`](@ref)).
Requires "POS", "VEL", and "MASS" to be available.
"""
function translate_to_center_of_velocity(p::Particles, radius::Number; kwargs...)
    v⃗ = center_of_velocity(p, radius; kwargs...)

    return translate(p, -v⃗, :vel)
end

function translate_to_center_of_velocity!(p::Particles, radius::Number; kwargs...)
    v⃗ = center_of_velocity(p, radius; kwargs...)

    return translate!(p, -v⃗, :vel)
end

"""
    translate_to_center_of_velocity(g::AbstractGalaxy, radius::Number, ptype::Symbol=:stars; kwargs...)
    translate_to_center_of_velocity!(g::AbstractGalaxy, radius::Number, ptype::Symbol=:stars; kwargs...)

Translates the particles to their center of velocity of the particles (see [`center_of_velocity`](@ref)).
Requires "POS", "VEL", and "MASS" of the particle type to be available.
"""
function translate_to_center_of_velocity(g::AbstractGalaxy, radius::Number, ptype::Symbol=:stars; kwargs...)
    v⃗ = center_of_velocity(g, radius, ptype; kwargs...)

    return translate(g, -v⃗, :vel)
end

function translate_to_center_of_velocity!(g::AbstractGalaxy, radius::Number, ptype::Symbol=:stars; kwargs...)
    v⃗ = center_of_velocity(g, radius, ptype; kwargs...)

    return translate!(g, -v⃗, :vel)
end


text:
#--------------------------------------------------------------------------
#  Example.jl
#
#  Test program for the functions in ComItAlg.jl, DiscAlg, and SimpRulT.jl.
#  The program also uses VAR1SimPs to calculate impulse response functions.
#
#  The program uses a very simplified version of the model in
#  Fuhrer, J. C. (1997), "Inflation/Output Variance Trade-Offs and Optimal
#  Monetary Policy," JMCB, 29, 214-234.
#
#
#
#  This example file and the procedures were created by translating MatLab code
#  to Julia.
#
#
#  Paul Soderlind (Paul.Soderlind@unisg.ch), 20 June 2000, to Julia Jan 2016
#----------------------------------------------------------------------------

using Dates, LinearAlgebra
include("jlFiles/Fuhrer1.jl")
include("jlFiles/VAR1SimPs.jl")
include("jlFiles/SimpRulT.jl")
include("jlFiles/ComitAlg.jl")
include("jlFiles/DiscAlg.jl")
include("jlFiles/printmat.jl")
#-----------------------------------------------------------------------

                               #simplified Fuhrer model, parameter values
(a1,ap,Varey)  = (0.85,-0.41,0.84^2)   #IS equation
(w,gamm,Varep) = (0.65,0.002,0.19^2)   #contracting equation
Dbig = 40                              #arbitrage condition
(qy,qpi,qf,bet) = (1,1,0.5,0.99)       #Loss function

                                #Parameters -> system matrices
(A,B,K,Q,U,R) = Fuhrer1(a1,ap,Varey,w,gamm,Varep,Dbig,qy,qpi,qf,bet)
(n1,n2,K)     = (3,2,1)

T   = 12                          #periods to simulate in VAR
Shock0 = [sqrt(Varep);zeros(2)]   #Impulse response wrt ep
#-----------------------------------------------------------------------

                         #SIMPLE RULE, u(t)=-Fx(t)
(Fy,Fpi) = (0.5,1.2)     #Simple decision rule: i(t) = = Fy*y(t) + Fpi*[p(t)-p(t-1)]
F   = [0  -Fy  -Fpi*4*(1-w)  0  -Fpi*4*w ]

x10 = zeros(n1)                      #initial state vector
SigmaXX = diagm(0=>[Varep,Varey,0])  #covariance matrix of shocks
(M_Simp,C_Simp,J0) = SimpRulT(A,B,Q,R,U,bet,n1,n2,F,SigmaXX,x10,1.0)

x1 = VAR1SimPs(M_Simp,Shock0,T)        #VAR of x1(t)
x2 = x1*C_Simp'                        #x2(t)
x  = [x1 x2]                           #x(t) = [x1(t),x2(t)]
uu = -x*F'                             #u(t)
ypii_Simp = [x[:,2]  (4*(1-w)*x[:,3]+4*w*x[:,5]) uu]

println("\nSimple rule: impulse response to a one std of price shock (y,pi,i)")
printTable(ypii_Simp,["y","π","i"],string.(1:T),cell00="period")
#---------------------------------------

#plotting, comment out this if PyPlot is not installed
using PyPlot
#PyPlot.svg(true)           #for ipynb notebooks
figure()
  plot(1:T,ypii_Simp)
  title("Simple rule: impulse response to a one std of price shock")
  legend(["Output","Inflation","Short interest rate"])
  #display(gcf())            #uncomment in Atom/Juno
#-----------------------------------------------------------------------

                         #COMMITMENT
(M_Commit,C_Commit) = ComItAlg(A,B,Q,R,U,bet,n1,n2,1.0)

k      = VAR1SimPs(M_Commit,[Shock0;zeros(2)],T) #VAR of x1(t),p2(t)
lambda = k*C_Commit'                             #x2(t),u(t),p1(t)
uu     = lambda[:,n2+1:n2+1]                     #u(t)
x      = [k[:,1:n1] lambda[:,1:n2]]              #x(t) = [x1(t),x2(t)]
ypii_Commit = [x[:,2] (4*(1-w)*x[:,3]+4*w*x[:,5]) uu]

println("\nCommitment: impulse response to a one std of price shock (y,pi,i)")
printTable(ypii_Commit,["y","π","i"],string.(1:T),cell00="period")
#---------------------------------------

#plotting, comment out this if PyPlot is not installed
figure()
  plot(1:T,ypii_Commit)
  title("Commitment: impulse response to a one std of price shock")
  legend(["Output","Inflation","Short interest rate"])
  #display(gcf())            #uncomment in Atom/Juno
#-----------------------------------------------------------------------

                         #DISCRETION
println("\nPlease wait, the discretionary case takes some time to solve")
(M_Disc,C_Disc,V,F) = DiscAlg(A,B,Q,R,U,bet,n1,n2,Matrix(1.0I,n1,n1),zeros(n2,n1),
                               [1e-1;1e-5],0,1,0,1,0,1e+4)

x1 = VAR1SimPs(M_Disc,Shock0,T)        #VAR of x1(t)
x2 = x1*C_Disc'                        #x2(t)
x  = [x1 x2]                           #x(t) = [x1(t),x2(t)]
uu = -x1*F'                            #u(t)
ypii_Disc = [x[:,2] (4*(1-w)*x[:,3]+4*w*x[:,5]) uu]

println("\nDiscretion: impulse response to a one std of price shock (y,pi,i)")
printTable(ypii_Disc,["y","π","i"],string.(1:T),cell00="period")
#---------------------------------------

#plotting, comment out this if PyPlot is not installed
figure()
  plot(1:T,ypii_Disc)
  title("Discretion: impulse response to a one std of price shock")
  legend(["Output","Inflation","Short interest rate"])
  #display(gcf())            #uncomment in Atom/Juno
#-----------------------------------------------------------------------

max_stars_repo_name: PaulSoderlind/PaulSoderlindCode

max_issues_repo_name: PaulSoderlind/PaulSoderlindCode

max_forks_repo_name: PaulSoderlind/PaulSoderlindCode


text:
using Flux
using Flux: train!, reset!
using Flux.Losses: mse

m = Chain(
  RNN(2, 5),
  Dense(5, 1))

function loss(x, y)
  reset!(m)
  sum(mse(m(xi), yi) for (xi, yi) in zip(x, y))
end

seq_init = [rand(Float32, 2)]
seq_1 = [rand(Float32, 2) for i = 1:3]
seq_2 = [rand(Float32, 2) for i = 1:3]

y1 = [rand(Float32, 1) for i = 1:3]
y2 = [rand(Float32, 1) for i = 1:3]

X = [seq_1, seq_2]
Y = [y1, y2]
data = zip(X, Y)

ps = params(m)
opt = ADAM(1e-3)

@show(loss(seq_1, y1))
for epoch in 1:20
  train!(loss, ps, data, opt)
  @show(loss(seq_1, y1))
end