text:
"""
This module defines a type called `Param` that is to be constructed
before performing
and kind of inversion.
The following functionality is currently added in this module:
* a simple RTM
* full-waveform inversion
* source and receiver filter inversion
Once a `Param` object is constructed, the following routines update certain fields
of the `Param` after performing the inversion.
* `xfwi!` : updates the initial subsurface models in `Param` using its observed data
* `wfwi!` : updates the source and receiver filters in `Param` using its observed data
"""
module FWI

using Misfits
using TimerOutputs
using LinearMaps
using LinearAlgebra
using Ipopt

import GeoPhyInv.Interpolation
import GeoPhyInv.Models
import GeoPhyInv.Acquisition
import GeoPhyInv.Data
import GeoPhyInv.Coupling
import GeoPhyInv.Fdtd
using Optim, LineSearches
using DistributedArrays
using Calculus
using Random

include("X.jl")

global const to = TimerOutput(); # create a timer object

# modeling types
struct ModFdtd end

struct ModFdtdBorn end

struct ModFdtdHBorn end

# fields
struct P end
struct Vx end
struct Vz end

"""
FWI Parameters

# Fields

* `mgrid::Vector{StepRangeLen}` : modelling grid
* `igrid::Vector{StepRangeLen}` : inversion grid
* `acqsrc::Acquisition.Src` : base source wavelet for modelling data
* `acqgeom::Acquisition.Geom` : acquisition geometry
* `tgrid::StepRangeLen` :
* `attrib_mod`
* `model_obs` : model used for generating observed data
* `model0` : background velocity model (only used during Born modeling and inversion)
* `parameterization` : a vector of Symbols specifying parameterization of the inversion vector
* `verbose` : print level on STOUT
* `attrib` : synthetic or real

TODO: add an extra attribute for coupling functions inversion and modelling
"""
mutable struct Param{Tmod, Tdatamisfit}
	"model inversion variable"
	mx::X{Float64,1}
	mxm::X{Float64,1}
	"forward modelling parameters for"
	paf::Tmod
	"base source wavelet"
	acqsrc::Acquisition.Src
	"adjoint source functions"
	adjsrc::Acquisition.Src
	"acquisition geometry"
	acqgeom::Acquisition.Geom
	"acquisition geometry for adjoint propagation"
	adjacqgeom::Acquisition.Geom
	"Seismic model on modeling grid"
	modm::Models.Seismic
	"background Seismic model"
	modm0::Models.Seismic
	"Seismic on inversion grid"
	modi::Models.Seismic
	"initialize on this model"
	mod_initial::Models.Seismic
	"prior model on the inversion grid"
	priori::Models.Seismic
	"gradient Seismic model on modeling grid"
	gmodm::Models.Seismic
	"gradient Seismic on inversion grid"
	gmodi::Models.Seismic
	"parameterization"
	parameterization::Vector{Symbol}
	"model preconditioning"
	mprecon::AbstractArray{Float64,2}
	"compute data misfit"
	paTD::Tdatamisfit
	paminterp::Interpolation.Kernel_2D{Float64}
	verbose::Bool
	"synthetic or field inversion"
	attrib::Symbol
end

Base.print(pa::Param)=nothing
Base.show(pa::Param)=nothing
Base.display(pa::Param)=nothing

# add data inverse covariance matrix here later
struct LS end # just cls inversion

struct LS_prior # cls inversion including prior
	# add data inverse covariance matrix here later, instead of just Float64
	pdgls::Float64
	# pdgls
	pmgls::Misfits.P_gls{Float64}
end
LS_prior(α1::Float64, Q::LinearMap)=LS_prior(α1, Misfits.P_gls(Q))

"""
this method constructs prior term with Q=α*I
* `ninv` : number of inversion variables, use `xfwi_ninv`
* `α` : scalar
"""
function LS_prior(ninv::Int, α=[1.0, 0.5])
	Q=LinearMap(
	     (y,x)->LinearAlgebra.mul!(y,x,α[2]),
	     (y,x)->LinearAlgebra.mul!(y,x,α[2]),
		      ninv, ninv, ismutating=true, issymmetric=true)
	return LS_prior(α[1],Misfits.P_gls(Q))
end

struct Migr end # returns first gradient

struct Migr_fd end # computes first gradient using FD

include("core.jl")
include("prior.jl")
include("xfwi.jl")
include("xwfwi.jl")

"""
Convert the data `TD` to `Src` after time reversal.
"""
function update_adjsrc!(adjsrc, δdat::Data.TD, adjacqgeom)
	(adjsrc.fields != δdat.fields) && error("dissimilar fields")
	nt=length(δdat.tgrid)
	for i in 1:adjacqgeom.nss
		for (j,field) in enumerate(δdat.fields)
			wav=adjsrc.wav[i,j]
			dat=δdat.d[i,j]
			for ir in 1:δdat.acqgeom.nr[i]
				for it in 1:nt
					@inbounds wav[it,ir]=value_adjsrc(dat[nt-it+1,ir],eval(field)())
				end
			end
		end
	end
	return nothing
end

value_adjsrc(s, ::P) = s
value_adjsrc(s, ::Vx) = -1.0*s # see adj tests
value_adjsrc(s, ::Vz) = -1.0*s

function generate_adjsrc(fields, tgrid, adjacqgeom)
	adjsrc=Acquisition.Src_zeros(adjacqgeom, fields, tgrid)
	return adjsrc
end

"""
Constructor for `Param`

# Arguments

* `acqsrc::Acquisition.Src` : source time functions
* `acqgeom::Acquisition.Geom` : acquisition geometry
* `tgrid::StepRangeLen` : modelling time grid
* `attrib_mod::Union{ModFdtd, ModFdtdBorn, ModFdtdHBorn}` : modelling attribute
* `modm::Models.Seismic` : seismic model on modelling mesh

# Optional Arguments
* `tgrid_obs::StepRangeLen` : time grid for observed data

* `igrid::Vector{StepRangeLen}=modm.mgrid` : inversion grid if different from the modelling grid, i.e., `modm.mgrid`
* `igrid_interp_scheme` : interpolation scheme
  * `=:B1` linear
  * `=:B2` second order
* `mprecon_factor::Float64=1` : factor to control model preconditioner, always greater than 1
  * `=1` means the preconditioning is switched off
  * `>1` means the preconditioning is switched on, larger the mprecon_factor, stronger the applied preconditioner becomes
* `dobs::Data.TD` : observed data
* `dprecon::Data.TD=Data.TD_ones(1,dobs.tgrid,dobs.acqgeom)` : data preconditioning, defaults to one
* `tlagssf_fracs=0.0` : maximum lagtime of unknown source filter
* `tlagrf_fracs=0.0` : maximum lagtime of unknown receiver filter
* `acqsrc_obs::Acquisition.Src=acqsrc` : source wavelets to generate *observed data*; can be different from `acqsrc`
* `modm_obs::Models.Seismic=modm` : actual seismic model to generate *observed data*
* `modm0::Models.Seismic=fill!(modm,0.0)` : background seismic model for Born modelling and inversion
* `parameterization::Vector{Symbol}` : subsurface parameterization
  * `=[:χvp, :χρ, :null]` for velocity and density inversion
  * `=[:null, :χρ, :null]` for density-only inversion
  * `=[:χvp, :null, :null]` for velocity-only inversion
* `verbose::Bool=false` : print level on STDOUT during inversion
* `attrib::Symbol=:synthetic` : an attribute to control class
  * `=:synthetic` synthetic data inversion
  * `=:field` field data inversion
* nworker : number of workers (input nothing to use all available)
"""
function Param(
	       acqsrc::Acquisition.Src,
	       acqgeom::Acquisition.Geom,
	       tgrid::StepRangeLen,
	       attrib_mod::Union{ModFdtd, ModFdtdBorn, ModFdtdHBorn},
	       modm::Models.Seismic;
	       # other optional
	       tgrid_obs::StepRangeLen=deepcopy(tgrid),
	       rfields=[:P],
	       igrid=nothing,
	       igrid_interp_scheme::Symbol=:B2,
	       mprecon_factor::Float64=1.0,
	       dobs::Data.TD=Data.TD_zeros(rfields,tgrid_obs,acqgeom),
	       dprecon=nothing,
	       tlagssf_fracs=0.0,
	       tlagrf_fracs=0.0,
	       acqsrc_obs::Acquisition.Src=deepcopy(acqsrc),
	       modm_obs::Models.Seismic=modm,
	       modm0=nothing,
	       parameterization::Vector{Symbol}=[:χvp, :χρ, :null],
	       verbose::Bool=false,
	       attrib::Symbol=:synthetic,
	       nworker=nothing,
	       )

	# make a copy of initial model
	modm=deepcopy(modm)

	if(igrid===nothing)
		igrid=deepcopy(modm.mgrid)
	end
	# igrid has to truncated because the gradient evaluation
	# is inaccurate on boundaries
	mg=modm.mgrid
	amin=max(igrid[1][1],mg[1][3])
	amax=min(igrid[1][end],mg[1][end-2])
	if(length(igrid[1])==1)
		grid1=range(amin,step=min(step(igrid[1]), mg[1][end-2]-amin),length=1)
	else
		grid1=range(amin,stop=amax,length=length(igrid[1]))
	end
	amin=max(mg[2][3],igrid[2][1])
	amax=min(igrid[2][end],mg[2][end-2])
	if(length(igrid[2])==1)
		grid2=range(amin,step=min(step(igrid[2]), mg[2][end-2]-amin),length=1)
	else
		grid2=range(amin,stop=amax,length=length(igrid[2]))
	end
	igrid=[grid1, grid2]

	# create modi according to igrid and interpolation of modm
	modi = Models.Seismic_zeros(igrid);
	# update bounds
	Models.adjust_bounds!(modi, modm)
	Models.interp_spray!(modm, modi, :interp)

	# save the initial model, as modi will be changed during the iterations
	mod_initial=deepcopy(modi)

	# allocate prior inputs
	priori=similar(modi)

	# use default background model modm0
	if(modm0 === nothing)
		modm0=deepcopy(modm) # some dummy
		fill!(modm0, 0.0)
	end
	isequal(modm0, modm_obs) && (@warn "modm0 == modm_obs")

	if(attrib == :synthetic)
		isequal(modm, modm_obs) && error("initial model same as actual model, zero misfit?")
	end

	# acqgeom geometry for adjoint propagation
	adjacqgeom = AdjGeom(acqgeom)

	# generate adjoint sources
	adjsrc=generate_adjsrc(rfields, tgrid, adjacqgeom)

	# generating forward and adjoint modelling engines
	# to generate modelled data, border values, etc.
	# most of the parameters given to this are dummy
	paf=Fdtd.Param(npw=2, model=modm, born_flag=true,
		acqgeom=[acqgeom, adjacqgeom], acqsrc=[acqsrc, adjsrc],
		rfields=rfields,
		sflags=[3, 2], rflags=[1, 1],
		backprop_flag=1,
		tgridmod=tgrid, gmodel_flag=true, verbose=verbose, illum_flag=false, nworker=nworker)

	gmodm=similar(modm)
	gmodi=similar(modi)

	#println("added fake precon")
	#dprecon=deepcopy(dobs)
	#fill!(dprecon, 1.0)
	#Data.taper!(dprecon, 20.)

	# check dprecon
	if(!(dprecon===nothing))
		(!(isapprox(dprecon,dobs))) && error("invalid dprecon used")
	end

	# create Parameters for data misfit
	#coup=Coupling.TD_delta(dobs.tgrid, tlagssf_fracs, tlagrf_fracs, rfields, acqgeom)
	# choose the data misfit
	#	(iszero(tlagssf_fracs))
 	# tlagssf_fracs==[0.0] | tlagssf_fracs=[0.0])
	# paTD=Data.P_misfit(Data.TD_zeros(rfields,tgrid,acqgeom),dobs,w=dprecon,coup=coup, func_attrib=optims[1]);

	if(iszero(dobs))
		Random.randn!(dobs) # dummy dobs, update later
	end

	paTD=Data.P_misfit(Data.TD_zeros(rfields,tgrid,acqgeom),dobs,w=dprecon);

	paminterp=Interpolation.Kernel([modi.mgrid[2], modi.mgrid[1]], [modm.mgrid[2], modm.mgrid[1]], igrid_interp_scheme)

	mx=X(prod(length.(modi.mgrid))*count(parameterization.≠:null),2)
	mxm=X(prod(length.(modm.mgrid))*count(parameterization.≠:null),2)

	# put modm as a vector, according to parameterization in mxm.x
	Models.Seismic_get!(mxm.x,modm,parameterization)

	pa = Param(
	     mx, mxm,
	     paf,
	     deepcopy(acqsrc),
	     Acquisition.Src_zeros(adjacqgeom, rfields, tgrid),
	     deepcopy(acqgeom),
	     adjacqgeom,
	     deepcopy(modm), deepcopy(modm0), modi, mod_initial,
	     priori,
	     gmodm,gmodi,
	     parameterization,
	     zeros(2,2), # dummy, update mprecon later
	     paTD,
	     paminterp,
	     verbose, attrib)

	# generate observed data if attrib is synthetic
	if((attrib == :synthetic))
		update_observed_data!(pa, modm_obs, attrib_mod)
	end

	iszero(dobs) && ((attrib == :real) ? error("input observed data for real data inversion") : error("problem generating synthetic observed data"))

	# generate modelled data
	F!(pa, nothing, attrib_mod)

#	build_mprecon!(pa, Array(pa.paf.c.illum_stack), mprecon_factor)
	pa.paf.c.illum_flag=false # switch off illum flag for speed

	# default weights are 1.0
	fill!(pa.mx.w, 1.0)

	# update priori in pa to defaults
	update_prior!(pa)

	return pa
end

"""
update the *synthetic* observed data in `Param`
* allocated memory, don't use in inner loops
"""
function update_observed_data!(pa::Param, modm_obs, attrib_mod=ModFdtd())
	# save modm of pa to put it back later
	modm_copy=deepcopy(pa.modm)

	# change modm in pa to actual model
	copyto!(pa.modm, modm_obs)

	# update models in the forward engine and do modelling
	# ModFdtdBorn: modm0 will be used as a background model for Born modelling and modm will be perturbed model
	# ModFdtd: modm will be *the* model
	F!(pa, nothing, attrib_mod)

	# get observed data
	copyto!(pa.paTD.y, pa.paTD.x)

	# put back model and sources
	copyto!(pa.modm, modm_copy)

	Fdtd.initialize!(pa.paf.c)  # clear everything
end

"""
Return the number of inversion variables for FWI corresponding to `Param`.
This number of inversion variables depend on the size of inversion mesh.
"""
function xfwi_ninv(pa::Param)
	length(pa.mx.x)
end

"""
Return the number of inversion variables for source and receiver filter inversion
corresponding to `Param`.
This number depends on the maximum lagtimes of the filters.
"""
function wfwi_ninv(pa::Param)
	return  pa.paTD.coup.tgridssf.nx
end

"build a model precon"
function build_mprecon!(pa,illum::Array{Float64}, mprecon_factor=1.0)
	"illum cannot be zero when building mprecon"
	(any(illum .<= 0.0)) && error("illum cannot be negative or zero")
	(mprecon_factor < 1.0)  && error("invalid mprecon_factor")

	illumi = zeros(length(pa.modi.mgrid[1]), length(pa.modi.mgrid[2]))
	Interpolation.interp_spray!(illumi, illum, pa.paminterp, :spray)
	(any(illumi .<= 0.0)) && error("illumi cannot be negative or zero")

	# use log to fix scaling of illum
	# (note that illum is not the exact diagonal of the Hessian matrix, so I am doing tricks here)
	# larger the mprecon_factor, stronger the preconditioning
	# division by maximum is safe because illum is non zero
	minillumi = minimum(illumi)
	maxillumi = maximum(illumi)
	illumi -= minillumi; # remove mean
	(maxillumi ≠ 0.0) && (illumi ./= maxillumi) # normalize using maxillumi
	#illumi = 1.0 + log.(1. + (mprecon_factor-1) * illumi) # illumi now lies between 1.0 and some value depending on how large illumi is
	illumi = (1. + (mprecon_factor-1) * illumi) # illumi now lies between 1.0 and some value depending on how large illumi is

	illumi = vcat([vec(illumi) for i in 1:count(pa.parameterization.≠:null)]...)

	length(illumi) == xfwi_ninv(pa) ? nothing : error("something went wrong with creating mprecon")
	pa.mprecon = spdiagm(illumi,(0),xfwi_ninv(pa), xfwi_ninv(pa))
end

# convert x to modm by interpolation
# x is on the sparse grid, and is changed by parameterization
function x_to_modm!(pa, x)

	# put sparse x into modi (useless step, just for visualization)
	Models.Seismic_reparameterize!(pa.modi, x, pa.parameterization)

	# get x on dense grid
	Interpolation.interp_spray!(x, pa.mxm.x, pa.paminterp, :interp, count(pa.parameterization.≠:null))

	# put dense x into modm, according to parameterization
	Models.Seismic_reparameterize!(pa.modm,pa.mxm.x,pa.parameterization)
end

function δx_to_δmods!(pa, δx)
	# get x on dense grid
	Interpolation.interp_spray!(δx, pa.mxm.x, pa.paminterp, :interp, count(pa.parameterization.≠:null))

	# reparameterize accordingly to get [δKI, δρI]
	Models.pert_reparameterize!(pa.paf.c.δmod, pa.mxm.x, pa.modm0, pa.parameterization)

	# input [δKI, δρI] to the modeling engine
	Fdtd.update_δmods!(pa.paf.c, pa.paf.c.δmod)
end

"""
convert the gradient output from Fdtd to gx
"""
function spray_gradient!(gx,  pa, ::ModFdtd)

	# apply chain rule on gmodm to get gradient w.r.t. dense x
	Models.gradient_chainrule!(pa.mxm.gx, pa.paf.c.gradient, pa.modm, pa.parameterization)

	# spray gradient w.r.t. dense x to the space of sparse x
	Interpolation.interp_spray!(gx, pa.mxm.gx, pa.paminterp, :spray, count(pa.parameterization.≠:null))

	# visualize gmodi here?
	# visualize_gx!(pa.gmodm, pa.modm, pa.gmodi, pa.modi, gx, pa)
end
function spray_gradient!(gx, pa, ::ModFdtdBorn)

	# apply chain rule on gmodm to get gradient w.r.t. dense x (note that background model is used for Born modeling here)
	Models.gradient_chainrule!(pa.mxm.gx, pa.paf.c.gradient, pa.modm, pa.parameterization)

	# spray gradient w.r.t. dense x to the space of sparse x
	Interpolation.interp_spray!(gx, pa.mxm.gx, pa.paminterp, :spray, count(pa.parameterization.≠:null))

	# visualize gmodi here?
	# visualize_gx!(pa.gmodm, pa.modm, pa.gmodi, pa.modi, gx, pa)
end

"""
Return bound vectors for the `Seismic` model,
depeding on paramaterization
"""
function Seismic_xbound!(lower_x, upper_x, pa)
	nznx = prod(length.(pa.modi.mgrid))

	bound1 = similar(lower_x)
	# create a Seismic model with minimum possible values
	modbound = deepcopy(pa.modi);
	modbound.χvp = Models.χ(fill(modbound.vp0[1], size(modbound.χvp)), modbound.ref.vp)
	modbound.χvs = Models.χ(fill(modbound.vs0[1], size(modbound.χvs)), modbound.ref.vs)
	modbound.χρ = Models.χ(fill(modbound.ρ0[1], size(modbound.χρ)), modbound.ref.ρ)

	Models.Seismic_get!(bound1,modbound,pa.parameterization)

	bound2=similar(upper_x)
	# create a Seismic model with maximum possible values
	modbound.χvp = Models.χ(fill(modbound.vp0[2], size(modbound.χvp)), modbound.ref.vp)
	modbound.χvs = Models.χ(fill(modbound.vs0[2], size(modbound.χvs)), modbound.ref.vs)
	modbound.χρ = Models.χ(fill(modbound.ρ0[2], size(modbound.χρ)), modbound.ref.ρ)

	Models.Seismic_get!(bound2,modbound,pa.parameterization)

	# this sorting operation is important because
	lower_x[:] = min.(bound1, bound2)
	upper_x[:] = max.(bound1, bound2)

	return modbound
end

# gmodi, gmodm <-- gx (just for visualizing gradient)
function visualize_gx!(gmodm, modm, gmodi, modi, gx, pa)
	# chain rule depending on re-parameterization
	Models.Seismic_chainrule!(gmodi, modi, gx, pa.parameterization, 1)

	# get gradient after interpolation (just for visualization, not exact)
	Models.interp_spray!(gmodi, gmodm, :interp, pa=pa.paminterp)
end

"""
Modify the input acquisition geometry
such that the adjoint source time functions can
be propagated from the receiver positions.
The number of supersources will remain the same.
All the recievers will be fired as simultaneous sources.
"""
function AdjGeom(geomin::Acquisition.Geom)
	geomout = deepcopy(geomin);
	geomout.sx = geomin.rx; geomout.sz = geomin.rz;
	geomout.ns = geomin.nr;

#	geomout.rx = geomin.sx; geomout.rz = geomin.sz;
#	geomout.nr = geomin.ns;

	geomout.rx = geomin.rx; geomout.rz = geomin.rz;
	geomout.nr = geomin.nr;

	return geomout
end

"""
Return functional and gradient of the LS objective
"""
function func_grad_Coupling!(storage, x::Vector{Float64},pa::Param)

	pa.verbose ? println("computing gradient...") : nothing

	# x to w
	Coupling_x!(x, pa, -1)

	if(storage === nothing)
		# compute misfit
		f = Data.error!(pa.paTD)
		return f
	else
		f = Data.error!(pa.paTD, :dJssf)
		Coupling_gx!(storage, pa.paTD.dJssf)
		return storage
	end
end

"""
Convert coupling functions to x and vice versa
"""
function Coupling_x!(x::Vector{Float64},
		   pa::Param,
		   flag::Int64)
	if(flag ==1) # convert w to x
		iss=1 # take only first source
		for ifield=1:length(pa.paTD.y.fields)
			w=pa.paTD.coup.ssf[iss,ifield]
			for i in eachindex(x)
				x[i]=w[i] # take only first source
			end
		end
	elseif(flag == -1) # convert x to w
		for iss=1:pa.paTD.y.acqgeom.nss, ifield=1:length(pa.paTD.y.fields)
			w=pa.paTD.coup.ssf[iss,ifield]
			for i in eachindex(x)
				w[i]=x[i]
			end
		end
	else
		error("invalid flag")
	end
end

"""
Convert dJssf to gradient vector
"""
function Coupling_gx!(gx, dJssf)
	gx[:]=0.0
	for ifield=1:size(dJssf,2), iss=1:size(dJssf,1)
		dwav=dJssf[iss,ifield]
		for i in eachindex(gx)
			gx[i]+=dwav[i]
		end
	end
end

end # module


text:
using StaticArrays

#Quadrature rulesfrom Shunn&Ham Paper

const tetrahedron1a = [
    SVector{4}([0.2500000000000000  0.2500000000000000  0.2500000000000000  0.2500000000000000])
]

const tetrahedron1w = [
  1.0,
]

const tetrahedron4a = [
    SVector{4}([0.5854101966249680  0.1381966011250110  0.1381966011250110  0.1381966011250110]),
    SVector{4}([0.1381966011250110  0.5854101966249680  0.1381966011250110  0.1381966011250110]),
    SVector{4}([0.1381966011250110  0.1381966011250110  0.5854101966249680  0.1381966011250110]),
    SVector{4}([0.1381966011250110  0.1381966011250110  0.1381966011250110  0.5854101966249680])
]

const tetrahedron4w = [
  0.2500000000000000,
  0.2500000000000000,
  0.2500000000000000,
  0.2500000000000000
]

const tetrahedron10a = [
    SVector{4}([0.7784952948213300  0.0738349017262234  0.0738349017262234  0.0738349017262234 ]),
    SVector{4}([0.0738349017262234  0.7784952948213300  0.0738349017262234  0.0738349017262234 ]),
    SVector{4}([0.0738349017262234  0.0738349017262234  0.7784952948213300  0.0738349017262234 ]),
    SVector{4}([0.0738349017262234  0.0738349017262234  0.0738349017262234  0.7784952948213300 ]),
    SVector{4}([0.4062443438840510  0.4062443438840510  0.0937556561159491  0.0937556561159491 ]),
    SVector{4}([0.4062443438840510  0.0937556561159491  0.4062443438840510  0.0937556561159491 ]),
    SVector{4}([0.4062443438840510  0.0937556561159491  0.0937556561159491  0.4062443438840510 ]),
    SVector{4}([0.0937556561159491  0.4062443438840510  0.4062443438840510  0.0937556561159491 ]),
    SVector{4}([0.0937556561159491  0.4062443438840510  0.0937556561159491  0.4062443438840510 ]),
    SVector{4}([0.0937556561159491  0.0937556561159491  0.4062443438840510  0.4062443438840510 ])
]

const tetrahedron10w = [
    0.0476331348432089,
    0.0476331348432089,
    0.0476331348432089,
    0.0476331348432089,
    0.1349112434378610,
    0.1349112434378610,
    0.1349112434378610,
    0.1349112434378610,
    0.1349112434378610,
    0.1349112434378610
]

const tetrahedron20a = [
    SVector{4}([ 0.9029422158182680  0.0323525947272439  0.0323525947272439  0.0323525947272439]),
    SVector{4}([ 0.0323525947272439  0.9029422158182680  0.0323525947272439  0.0323525947272439]),
    SVector{4}([ 0.0323525947272439  0.0323525947272439  0.9029422158182680  0.0323525947272439]),
    SVector{4}([ 0.0323525947272439  0.0323525947272439  0.0323525947272439  0.9029422158182680]),
    SVector{4}([ 0.2626825838877790  0.6165965330619370  0.0603604415251421  0.0603604415251421]),
    SVector{4}([ 0.6165965330619370  0.2626825838877790  0.0603604415251421  0.0603604415251421]),
    SVector{4}([ 0.2626825838877790  0.0603604415251421  0.6165965330619370  0.0603604415251421]),
    SVector{4}([ 0.6165965330619370  0.0603604415251421  0.2626825838877790  0.0603604415251421]),
    SVector{4}([ 0.2626825838877790  0.0603604415251421  0.0603604415251421  0.6165965330619370]),
    SVector{4}([ 0.6165965330619370  0.0603604415251421  0.0603604415251421  0.2626825838877790]),
    SVector{4}([ 0.0603604415251421  0.2626825838877790  0.6165965330619370  0.0603604415251421]),
    SVector{4}([ 0.0603604415251421  0.6165965330619370  0.2626825838877790  0.0603604415251421]),
    SVector{4}([ 0.0603604415251421  0.2626825838877790  0.0603604415251421  0.6165965330619370]),
    SVector{4}([ 0.0603604415251421  0.6165965330619370  0.0603604415251421  0.2626825838877790]),
    SVector{4}([ 0.0603604415251421  0.0603604415251421  0.2626825838877790  0.6165965330619370]),
    SVector{4}([ 0.0603604415251421  0.0603604415251421  0.6165965330619370  0.2626825838877790]),
    SVector{4}([ 0.3097693042728620  0.3097693042728620  0.3097693042728620  0.0706920871814129]),
    SVector{4}([ 0.3097693042728620  0.3097693042728620  0.0706920871814129  0.3097693042728620]),
    SVector{4}([ 0.3097693042728620  0.0706920871814129  0.3097693042728620  0.3097693042728620]),
    SVector{4}([ 0.0706920871814129  0.3097693042728620  0.3097693042728620  0.3097693042728620])
]

const tetrahedron20w = [
    0.0070670747944695,
    0.0070670747944695,
    0.0070670747944695,
    0.0070670747944695,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.0469986689718877,
    0.1019369182898680,
    0.1019369182898680,
    0.1019369182898680,
    0.1019369182898680
]

const tetrahedron35a = [
   SVector{4}([0.9197896733368800  0.0267367755543735  0.0267367755543735  0.0267367755543735]),
   SVector{4}([0.0267367755543735  0.9197896733368800  0.0267367755543735  0.0267367755543735]),
   SVector{4}([0.0267367755543735  0.0267367755543735  0.9197896733368800  0.0267367755543735]),
   SVector{4}([0.0267367755543735  0.0267367755543735  0.0267367755543735  0.9197896733368800]),
   SVector{4}([0.1740356302468940  0.7477598884818090  0.0391022406356488  0.0391022406356488]),
   SVector{4}([0.7477598884818090  0.1740356302468940  0.0391022406356488  0.0391022406356488]),
   SVector{4}([0.1740356302468940  0.0391022406356488  0.7477598884818090  0.0391022406356488]),
   SVector{4}([0.7477598884818090  0.0391022406356488  0.1740356302468940  0.0391022406356488]),
   SVector{4}([0.1740356302468940  0.0391022406356488  0.0391022406356488  0.7477598884818090]),
   SVector{4}([0.7477598884818090  0.0391022406356488  0.0391022406356488  0.1740356302468940]),
   SVector{4}([0.0391022406356488  0.1740356302468940  0.7477598884818090  0.0391022406356488]),
   SVector{4}([0.0391022406356488  0.7477598884818090  0.1740356302468940  0.0391022406356488]),
   SVector{4}([0.0391022406356488  0.1740356302468940  0.0391022406356488  0.7477598884818090]),
   SVector{4}([0.0391022406356488  0.7477598884818090  0.0391022406356488  0.1740356302468940]),
   SVector{4}([0.0391022406356488  0.0391022406356488  0.1740356302468940  0.7477598884818090]),
   SVector{4}([0.0391022406356488  0.0391022406356488  0.7477598884818090  0.1740356302468940]),
   SVector{4}([0.4547545999844830  0.4547545999844830  0.0452454000155172  0.0452454000155172]),
   SVector{4}([0.4547545999844830  0.0452454000155172  0.4547545999844830  0.0452454000155172]),
   SVector{4}([0.4547545999844830  0.0452454000155172  0.0452454000155172  0.4547545999844830]),
   SVector{4}([0.0452454000155172  0.4547545999844830  0.4547545999844830  0.0452454000155172]),
   SVector{4}([0.0452454000155172  0.4547545999844830  0.0452454000155172  0.4547545999844830]),
   SVector{4}([0.0452454000155172  0.0452454000155172  0.4547545999844830  0.4547545999844830]),
   SVector{4}([0.5031186450145980  0.2232010379623150  0.2232010379623150  0.0504792790607720]),
   SVector{4}([0.2232010379623150  0.5031186450145980  0.2232010379623150  0.0504792790607720]),
   SVector{4}([0.2232010379623150  0.2232010379623150  0.5031186450145980  0.0504792790607720]),
   SVector{4}([0.5031186450145980  0.2232010379623150  0.0504792790607720  0.2232010379623150]),
   SVector{4}([0.2232010379623150  0.5031186450145980  0.0504792790607720  0.2232010379623150]),
   SVector{4}([0.2232010379623150  0.2232010379623150  0.0504792790607720  0.5031186450145980]),
   SVector{4}([0.5031186450145980  0.0504792790607720  0.2232010379623150  0.2232010379623150]),
   SVector{4}([0.2232010379623150  0.0504792790607720  0.5031186450145980  0.2232010379623150]),
   SVector{4}([0.2232010379623150  0.0504792790607720  0.2232010379623150  0.5031186450145980]),
   SVector{4}([0.0504792790607720  0.5031186450145980  0.2232010379623150  0.2232010379623150]),
   SVector{4}([0.0504792790607720  0.2232010379623150  0.5031186450145980  0.2232010379623150]),
   SVector{4}([0.0504792790607720  0.2232010379623150  0.2232010379623150  0.5031186450145980]),
   SVector{4}([0.2500000000000000  0.2500000000000000  0.2500000000000000  0.2500000000000000])
]

const tetrahedron35w = [
    0.0021900463965388,
    0.0021900463965388,
    0.0021900463965388,
    0.0021900463965388,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0143395670177665,
    0.0250305395686746,
    0.0250305395686746,
    0.0250305395686746,
    0.0250305395686746,
    0.0250305395686746,
    0.0250305395686746,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0479839333057554,
    0.0931745731195340
]

const tetrahedron56a = [
   SVector{4}([0.9551438045408220  0.0149520651530592  0.0149520651530592  0.0149520651530592]),
   SVector{4}([0.0149520651530592  0.9551438045408220  0.0149520651530592  0.0149520651530592]),
   SVector{4}([0.0149520651530592  0.0149520651530592  0.9551438045408220  0.0149520651530592]),
   SVector{4}([0.0149520651530592  0.0149520651530592  0.0149520651530592  0.9551438045408220]),
   SVector{4}([0.7799760084415400  0.1518319491659370  0.0340960211962615  0.0340960211962615]),
   SVector{4}([0.1518319491659370  0.7799760084415400  0.0340960211962615  0.0340960211962615]),
   SVector{4}([0.7799760084415400  0.0340960211962615  0.1518319491659370  0.0340960211962615]),
   SVector{4}([0.1518319491659370  0.0340960211962615  0.7799760084415400  0.0340960211962615]),
   SVector{4}([0.7799760084415400  0.0340960211962615  0.0340960211962615  0.1518319491659370]),
   SVector{4}([0.1518319491659370  0.0340960211962615  0.0340960211962615  0.7799760084415400]),
   SVector{4}([0.0340960211962615  0.7799760084415400  0.1518319491659370  0.0340960211962615]),
   SVector{4}([0.0340960211962615  0.1518319491659370  0.7799760084415400  0.0340960211962615]),
   SVector{4}([0.0340960211962615  0.7799760084415400  0.0340960211962615  0.1518319491659370]),
   SVector{4}([0.0340960211962615  0.1518319491659370  0.0340960211962615  0.7799760084415400]),
   SVector{4}([0.0340960211962615  0.0340960211962615  0.7799760084415400  0.1518319491659370]),
   SVector{4}([0.0340960211962615  0.0340960211962615  0.1518319491659370  0.7799760084415400]),
   SVector{4}([0.3549340560639790  0.5526556431060170  0.0462051504150017  0.0462051504150017]),
   SVector{4}([0.5526556431060170  0.3549340560639790  0.0462051504150017  0.0462051504150017]),
   SVector{4}([0.3549340560639790  0.0462051504150017  0.5526556431060170  0.0462051504150017]),
   SVector{4}([0.5526556431060170  0.0462051504150017  0.3549340560639790  0.0462051504150017]),
   SVector{4}([0.3549340560639790  0.0462051504150017  0.0462051504150017  0.5526556431060170]),
   SVector{4}([0.5526556431060170  0.0462051504150017  0.0462051504150017  0.3549340560639790]),
   SVector{4}([0.0462051504150017  0.3549340560639790  0.5526556431060170  0.0462051504150017]),
   SVector{4}([0.0462051504150017  0.5526556431060170  0.3549340560639790  0.0462051504150017]),
   SVector{4}([0.0462051504150017  0.3549340560639790  0.0462051504150017  0.5526556431060170]),
   SVector{4}([0.0462051504150017  0.5526556431060170  0.0462051504150017  0.3549340560639790]),
   SVector{4}([0.0462051504150017  0.0462051504150017  0.3549340560639790  0.5526556431060170]),
   SVector{4}([0.0462051504150017  0.0462051504150017  0.5526556431060170  0.3549340560639790]),
   SVector{4}([0.5381043228880020  0.2281904610687610  0.2281904610687610  0.0055147549744775]),
   SVector{4}([0.2281904610687610  0.5381043228880020  0.2281904610687610  0.0055147549744775]),
   SVector{4}([0.2281904610687610  0.2281904610687610  0.5381043228880020  0.0055147549744775]),
   SVector{4}([0.5381043228880020  0.2281904610687610  0.0055147549744775  0.2281904610687610]),
   SVector{4}([0.2281904610687610  0.5381043228880020  0.0055147549744775  0.2281904610687610]),
   SVector{4}([0.2281904610687610  0.2281904610687610  0.0055147549744775  0.5381043228880020]),
   SVector{4}([0.5381043228880020  0.0055147549744775  0.2281904610687610  0.2281904610687610]),
   SVector{4}([0.2281904610687610  0.0055147549744775  0.5381043228880020  0.2281904610687610]),
   SVector{4}([0.2281904610687610  0.0055147549744775  0.2281904610687610  0.5381043228880020]),
   SVector{4}([0.0055147549744775  0.5381043228880020  0.2281904610687610  0.2281904610687610]),
   SVector{4}([0.0055147549744775  0.2281904610687610  0.5381043228880020  0.2281904610687610]),
   SVector{4}([0.0055147549744775  0.2281904610687610  0.2281904610687610  0.5381043228880020]),
   SVector{4}([0.1961837595745600  0.3523052600879940  0.3523052600879940  0.0992057202494530]),
   SVector{4}([0.3523052600879940  0.1961837595745600  0.3523052600879940  0.0992057202494530]),
   SVector{4}([0.3523052600879940  0.3523052600879940  0.1961837595745600  0.0992057202494530]),
   SVector{4}([0.1961837595745600  0.3523052600879940  0.0992057202494530  0.3523052600879940]),
   SVector{4}([0.3523052600879940  0.1961837595745600  0.0992057202494530  0.3523052600879940]),
   SVector{4}([0.3523052600879940  0.3523052600879940  0.0992057202494530  0.1961837595745600]),
   SVector{4}([0.1961837595745600  0.0992057202494530  0.3523052600879940  0.3523052600879940]),
   SVector{4}([0.3523052600879940  0.0992057202494530  0.1961837595745600  0.3523052600879940]),
   SVector{4}([0.3523052600879940  0.0992057202494530  0.3523052600879940  0.1961837595745600]),
   SVector{4}([0.0992057202494530  0.1961837595745600  0.3523052600879940  0.3523052600879940]),
   SVector{4}([0.0992057202494530  0.3523052600879940  0.1961837595745600  0.3523052600879940]),
   SVector{4}([0.0992057202494530  0.3523052600879940  0.3523052600879940  0.1961837595745600]),
   SVector{4}([0.5965649956210170  0.1344783347929940  0.1344783347929940  0.1344783347929940]),
   SVector{4}([0.1344783347929940  0.5965649956210170  0.1344783347929940  0.1344783347929940]),
   SVector{4}([0.1344783347929940  0.1344783347929940  0.5965649956210170  0.1344783347929940]),
   SVector{4}([0.1344783347929940  0.1344783347929940  0.1344783347929940  0.5965649956210170])
]

const tetrahedron56w = [
    0.0010373112336140,
    0.0010373112336140,
    0.0010373112336140,
    0.0010373112336140,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0096016645399480,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0164493976798232,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0153747766513310,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0293520118375230,
    0.0366291366405108,
    0.0366291366405108,
    0.0366291366405108,
    0.0366291366405108
]

const tetrahedron85a = [
    SVector{4}([0.919364576755549   0.026878474414817   0.026878474414817   0.026878474414817]),
    SVector{4}([0.026878474414817   0.919364576755549   0.026878474414817   0.026878474414817]),
    SVector{4}([0.026878474414817   0.026878474414817   0.919364576755549   0.026878474414817]),
    SVector{4}([0.026878474414817   0.026878474414817   0.026878474414817   0.919364576755549]),
    SVector{4}([0.438577972589591   0.187140675803470   0.187140675803470   0.187140675803470]),
    SVector{4}([0.187140675803470   0.438577972589591   0.187140675803470   0.187140675803470]),
    SVector{4}([0.187140675803470   0.187140675803470   0.438577972589591   0.187140675803470]),
    SVector{4}([0.187140675803470   0.187140675803470   0.187140675803470   0.438577972589591]),
    SVector{4}([0.473575835127937   0.473575835127937   0.026424164872063   0.026424164872063]),
    SVector{4}([0.473575835127937   0.026424164872063   0.473575835127937   0.026424164872063]),
    SVector{4}([0.473575835127937   0.026424164872063   0.026424164872063   0.473575835127937]),
    SVector{4}([0.026424164872063   0.473575835127937   0.473575835127937   0.026424164872063]),
    SVector{4}([0.026424164872063   0.473575835127937   0.026424164872063   0.473575835127937]),
    SVector{4}([0.026424164872063   0.026424164872063   0.473575835127937   0.473575835127937]),
    SVector{4}([0.352045262027356   0.352045262027356   0.147954737972644   0.147954737972644]),
    SVector{4}([0.352045262027356   0.147954737972644   0.352045262027356   0.147954737972644]),
    SVector{4}([0.352045262027356   0.147954737972644   0.147954737972644   0.352045262027356]),
    SVector{4}([0.147954737972644   0.352045262027356   0.352045262027356   0.147954737972644]),
    SVector{4}([0.147954737972644   0.352045262027356   0.147954737972644   0.352045262027356]),
    SVector{4}([0.147954737972644   0.147954737972644   0.352045262027356   0.352045262027356]),
    SVector{4}([0.225783205866940   0.732309909692947   0.020953442220056   0.020953442220056]),
    SVector{4}([0.225783205866940   0.020953442220056   0.732309909692947   0.020953442220056]),
    SVector{4}([0.225783205866940   0.020953442220056   0.020953442220056   0.732309909692947]),
    SVector{4}([0.020953442220056   0.225783205866940   0.732309909692947   0.020953442220056]),
    SVector{4}([0.020953442220056   0.225783205866940   0.020953442220056   0.732309909692947]),
    SVector{4}([0.020953442220056   0.020953442220056   0.225783205866940   0.732309909692947]),
    SVector{4}([0.732309909692947   0.225783205866940   0.020953442220056   0.020953442220056]),
    SVector{4}([0.732309909692947   0.020953442220056   0.225783205866940   0.020953442220056]),
    SVector{4}([0.732309909692947   0.020953442220056   0.020953442220056   0.225783205866940]),
    SVector{4}([0.020953442220056   0.732309909692947   0.225783205866940   0.020953442220056]),
    SVector{4}([0.020953442220056   0.732309909692947   0.020953442220056   0.225783205866940]),
    SVector{4}([0.020953442220056   0.020953442220056   0.732309909692947   0.225783205866940]),
    SVector{4}([0.158462939666092   0.647557594086975   0.096989733123466   0.096989733123466]),
    SVector{4}([0.158462939666092   0.096989733123466   0.647557594086975   0.096989733123466]),
    SVector{4}([0.158462939666092   0.096989733123466   0.096989733123466   0.647557594086975]),
    SVector{4}([0.096989733123466   0.158462939666092   0.647557594086975   0.096989733123466]),
    SVector{4}([0.096989733123466   0.158462939666092   0.096989733123466   0.647557594086975]),
    SVector{4}([0.096989733123466   0.096989733123466   0.158462939666092   0.647557594086975]),
    SVector{4}([0.647557594086975   0.158462939666092   0.096989733123466   0.096989733123466]),
    SVector{4}([0.647557594086975   0.096989733123466   0.158462939666092   0.096989733123466]),
    SVector{4}([0.647557594086975   0.096989733123466   0.096989733123466   0.158462939666092]),
    SVector{4}([0.096989733123466   0.647557594086975   0.158462939666092   0.096989733123466]),
    SVector{4}([0.096989733123466   0.647557594086975   0.096989733123466   0.158462939666092]),
    SVector{4}([0.096989733123466   0.096989733123466   0.647557594086975   0.158462939666092]),
    SVector{4}([0.322111431830857   0.322111431830857   0.322111431830857   0.033665704507429]),
    SVector{4}([0.322111431830857   0.322111431830857   0.033665704507429   0.322111431830857]),
    SVector{4}([0.322111431830857   0.033665704507429   0.322111431830857   0.322111431830857]),
    SVector{4}([0.033665704507429   0.322111431830857   0.322111431830857   0.322111431830857]),
    SVector{4}([0.792939256469618   0.097608162890442   0.097608162890442   0.011844417749498]),
    SVector{4}([0.097608162890442   0.792939256469618   0.097608162890442   0.011844417749498]),
    SVector{4}([0.097608162890442   0.097608162890442   0.792939256469618   0.011844417749498]),
    SVector{4}([0.792939256469618   0.097608162890442   0.011844417749498   0.097608162890442]),
    SVector{4}([0.097608162890442   0.792939256469618   0.011844417749498   0.097608162890442]),
    SVector{4}([0.097608162890442   0.097608162890442   0.011844417749498   0.792939256469618]),
    SVector{4}([0.792939256469618   0.011844417749498   0.097608162890442   0.097608162890442]),
    SVector{4}([0.097608162890442   0.011844417749498   0.792939256469618   0.097608162890442]),
    SVector{4}([0.097608162890442   0.011844417749498   0.097608162890442   0.792939256469618]),
    SVector{4}([0.011844417749498   0.792939256469618   0.097608162890442   0.097608162890442]),
    SVector{4}([0.011844417749498   0.097608162890442   0.792939256469618   0.097608162890442]),
    SVector{4}([0.011844417749498   0.097608162890442   0.097608162890442   0.792939256469618]),
    SVector{4}([0.541184412800237   0.133558160703568   0.296501020543124   0.028756405953071]),
    SVector{4}([0.296501020543124   0.541184412800237   0.133558160703568   0.028756405953071]),
    SVector{4}([0.133558160703568   0.296501020543124   0.541184412800237   0.028756405953071]),
    SVector{4}([0.541184412800237   0.296501020543124   0.133558160703568   0.028756405953071]),
    SVector{4}([0.133558160703568   0.541184412800237   0.296501020543124   0.028756405953071]),
    SVector{4}([0.296501020543124   0.133558160703568   0.541184412800237   0.028756405953071]),
    SVector{4}([0.541184412800237   0.133558160703568   0.028756405953071   0.296501020543124]),
    SVector{4}([0.296501020543124   0.541184412800237   0.028756405953071   0.133558160703568]),
    SVector{4}([0.133558160703568   0.296501020543124   0.028756405953071   0.541184412800237]),
    SVector{4}([0.541184412800237   0.296501020543124   0.028756405953071   0.133558160703568]),
    SVector{4}([0.133558160703568   0.541184412800237   0.028756405953071   0.296501020543124]),
    SVector{4}([0.296501020543124   0.133558160703568   0.028756405953071   0.541184412800237]),
    SVector{4}([0.541184412800237   0.028756405953071   0.133558160703568   0.296501020543124]),
    SVector{4}([0.296501020543124   0.028756405953071   0.541184412800237   0.133558160703568]),
    SVector{4}([0.133558160703568   0.028756405953071   0.296501020543124   0.541184412800237]),
    SVector{4}([0.541184412800237   0.028756405953071   0.296501020543124   0.133558160703568]),
    SVector{4}([0.133558160703568   0.028756405953071   0.541184412800237   0.296501020543124]),
    SVector{4}([0.296501020543124   0.028756405953071   0.133558160703568   0.541184412800237]),
    SVector{4}([0.028756405953071   0.541184412800237   0.133558160703568   0.296501020543124]),
    SVector{4}([0.028756405953071   0.296501020543124   0.541184412800237   0.133558160703568]),
    SVector{4}([0.028756405953071   0.133558160703568   0.296501020543124   0.541184412800237]),
    SVector{4}([0.028756405953071   0.541184412800237   0.296501020543124   0.133558160703568]),
    SVector{4}([0.028756405953071   0.133558160703568   0.541184412800237   0.296501020543124]),
    SVector{4}([0.028756405953071   0.296501020543124   0.133558160703568   0.541184412800237])
]

const tetrahedron85w = [
    0.002144935144316,
    0.002144935144316,
    0.002144935144316,
    0.002144935144316,
    0.020826641690769,
    0.020826641690769,
    0.020826641690769,
    0.020826641690769,
    0.007210136064455,
    0.007210136064455,
    0.007210136064455,
    0.007210136064455,
    0.007210136064455,
    0.007210136064455,
    0.030798919159712,
    0.030798919159712,
    0.030798919159712,
    0.030798919159712,
    0.030798919159712,
    0.030798919159712,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.004357844813864,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.008593530677833,
    0.023000681669286,
    0.023000681669286,
    0.023000681669286,
    0.023000681669286,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.004863063904912,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259,
    0.015595140078259
]

const tetrahedronA = [
  tetrahedron1a,
  tetrahedron4a,
  tetrahedron10a,
  tetrahedron20a,
  tetrahedron35a,
  tetrahedron56a,
  tetrahedron85a
]

const tetrahedronW = [
  tetrahedron1w,
  tetrahedron4w,
  tetrahedron10w,
  tetrahedron20w,
  tetrahedron35w,
  tetrahedron56w,
  tetrahedron85w
]

ext: jl


text:
using ReverseDiff
using Distributions
include("char_loader.jl")

loader = CharLoader("../data/input.txt")

# number of unique chars in the dataset
BATCH_SIZE = 25
INPUT_SIZE = length(loader.chars)
OUTPUT_SIZE = length(loader.chars)
STATE_SIZE = 100

Wh = randn(T, state_size, state_size) * 0.01
Wx = randn(T, state_size, input_size) * 0.01
Wo = randn(T, output_size, state_size) * 0.01

ReverseDiff.@forward minus_log(x) = -log(x)
softmax(x) = (exp_x = exp.(x); exp_x ./ sum(exp_x, 1))
cross_entropy_loss(x, y) = mean(sum(y .* minus_log.(x), 1))

function train_rnn(cell::AbstractCell{T}, X, Y)
    loss = T(0)
    n = size(X, 2)
    ht = zeros(1, state_size)
    for t=1:n
        ot, ht = rnn_cell(Wx, Wh, Wo, X[t, :], ht)
        y = Y[t]
        loss += minus_log.(softmax(ot))[y]
    end
    return loss / n
end

function predict_neural_net(Wx, Wh, Wo, X)
    preds = Int[]
    T = size(X, 1)
    ht = zeros(1, state_size)
    for t=1:T
        xt = view(X, t, :)
        ot, ht = rnn_cell(Wx, Wh, Wo, xt, ht)
        push!(preds, indmax(ot))
    end
    return preds
end

function generate_text(T)
    id0 = rand(1:input_size, 1)[1]
    xt = zeros(1, input_size)
    xt[id0] = 1
    ht = zeros(1, state_size)
    ids = Int[id0]

    for t=1:T
        ot, ht = rnn_cell(Wx, Wh, Wo, xt, ht)
        probs = softmax(ot)
        ct = Categorical(probs[:])
        pred = rand(ct)
        xt[:] = 0.
        xt[pred] = 1.
        push!(ids, pred)
    end
    return ids
end

function train_model(iters, Wx, Wh, Wo; α=1e-2)
    for i=1:iters
        x, y = next_batch(loader, batch_size)
        dWx, dWh, dWo = nn_backward(Wx, Wh, Wo, x, y)
        # Vanilla SGD
        Wx += α * dWx
        Wh += α * dWh
        Wo += α * dWo

        if i % 100 == 0
            println("Loop iter $i")
            println("Generated text ...")
            ids = generate_text(200)
            for i in ids
                print(loader.id2char[i])
            end
            println("\n")
        end
    end
end

accuracy(ypred, ytrue) = mean(ypred .== ytrue)


text:
using FIEREE, PyCall, JSON, Dates

cloud_project = "ee-sandbox"
secret_key = "/home/Socrates/kmarkert/fier_testing/ee-sandbox.json"

# initialize Earth Engine
ee = pyimport("ee")
ee.Initialize()

session = FIEREE.get_session(secret_key)

# define parameters for space and time info
resolution = 0.005
start_time = "2015-01-01";
end_time = "2020-01-01";

kh = ee.FeatureCollection("USDOS/LSIB_SIMPLE/2017").filter(ee.Filter.eq("country_na","Cambodia"));

basins = ee.FeatureCollection("WWF/HydroSHEDS/v1/Basins/hybas_7").filterBounds(kh);
basin_codes = basins.aggregate_array("HYBAS_ID").getInfo()

cnt = 1
total = length(basin_codes)

for code in basin_codes
    now = Dates.now()
    println("$now -- ($cnt/$total): Processing basin code $code")
    basin = basins.filter(ee.Filter.eq("HYBAS_ID",code));

    box = FIEREE.geom_to_bbox(basin)

    domain = FIEREE.get_domain(box,resolution);
    println(domain)

    dates = FIEREE.get_s1_dates(domain,start_time,end_time);
    data_arr = FIEREE.get_s1_data(cloud_project,session,domain,start_time,end_time);

    spatial_modes,temporal_modes = FIEREE.reof(data_arr;n_modes=4);

    df = FIEREE.get_spt_table(basin);
    df_matched = FIEREE.match_dates(df,dates);

    results = FIEREE.find_fits(df_matched,spatial_modes,temporal_modes,data_arr);

    open("fier_results_$code.json","w") do io
        JSON.print(io,results,4)
    end

end


text:
module VSCodeDebugger

include("../../CodeTracking/src/CodeTracking.jl")

module JuliaInterpreter
    using ..CodeTracking

    include("../../JuliaInterpreter/src/packagedef.jl")
end

module DebugAdapter
    import ..JuliaInterpreter

    include("../../DebugAdapter/src/packagedef.jl")
end

end

max_stars_repo_name: fonsp/julia-vscode

max_issues_repo_name: fonsp/julia-vscode

max_forks_repo_name: fonsp/julia-vscode


text:
struct ParameterStateGradientWatcher
    p::Dict
    s::Dict
end

function capture_parameter_state_references(model, base_name=nothing; ps::Zygote.IdSet=Flux.params(model).params)
    parameter_references = Dict{String,AbstractArray}()
    state_references = Dict{String,AbstractArray}()

    fmap_custom(model, base_name) do x, k
        if x isa AbstractArray{<:Number}
            if x ∈ ps
                parameter_references[string(k)] = x
            else
                state_references[string(k)] = x
            end
        end
        return x
    end

    return parameter_references, state_references
end

function ParameterStateGradientWatcher(model, base_name::Symbol=Symbol(); ps::Zygote.IdSet=Flux.params(model).params)
    parameter_references, state_references = capture_parameter_state_references(model, base_name; ps=ps)
    return ParameterStateGradientWatcher(parameter_references, state_references)
end

function Base.log(lg::WandbLogger, ps::ParameterStateGradientWatcher, gs::Zygote.Grads; kwargs...)
    log_dict = Dict{String,Any}()
    for (pname, p) in ps.p
        log_dict["Parameters/" * pname] = Wandb.Histogram(p |> cpu)
        log_dict["Gradients/" * pname] = Wandb.Histogram(gs[p] |> cpu)
    end
    for (sname, s) in ps.s
        log_dict["States/" * sname] = Wandb.Histogram(s |> cpu)
    end
    return log(lg, log_dict; kwargs...)
end


text:
using KitBase, FluxReconstruction, Langevin, LinearAlgebra, OrdinaryDiffEq, OffsetArrays, Plots, JLD2
using ProgressMeter: @showprogress
using Base.Threads: @threads

begin
    set = Setup(
        "gas",
        "cylinder",
        "2d0f0v",
        "hll",
        "nothing",
        1, # species
        3, # order of accuracy
        "positivity", # limiter
        "euler",
        0.1, # cfl
        1.0, # time
    )
    ps = FRPSpace2D(0.0, 2.0, 100, 0.0, 1.0, 50, set.interpOrder-1, 1, 1)
    vs = nothing
    gas = Gas(
        1e-6,
        1.12, # Mach
        1.0,
        3.0, # K
        7/5,
        0.81,
        1.0,
        0.5,
    )
    ib = nothing
    ks = SolverSet(set, ps, vs, gas, ib)
end

begin
    uqMethod = "collocation"
    nr = 5
    nRec = 10
    opType = "uniform"
    parameter1 = 0.95
    parameter2 = 1.05

    uq = UQ1D(nr, nRec, parameter1, parameter2, opType, uqMethod)
end
l2 = [uq.t2Product[j-1, j-1] for j = 1:uq.nm+1]

cd(@__DIR__)
include("../ic.jl")
include("../../filter.jl")

function dudt!(du, u, p, t)
    du .= 0.0

    f, u_face, f_face, fx_interaction, fy_interaction, rhs1, rhs2,
    iJ, ll, lr, dhl, dhr, lpdm, γ, uq = p

    nx = size(u, 1) - 2
    ny = size(u, 2) - 2
    nsp = size(u, 3)
    nq = size(u, 6)

    @inbounds @threads for l = 1:nsp
        for k = 1:nsp, j in axes(f, 2), i in axes(f, 1)
            for n = 1:nq
                fg, gg = euler_flux(u[i, j, k, l, :, n], γ)
                for s = 1:4
                    f[i, j, k, l, s, n, :] .= iJ[i, j][k, l] * [fg[s], gg[s]]
                end
            end
        end
    end

    @inbounds @threads for n = 1:nq
        for m = 1:4, l = 1:nsp, j in axes(u_face, 2), i in axes(u_face, 1)
            u_face[i, j, 1, l, m, n] = dot(u[i, j, l, :, m, n], ll)
            u_face[i, j, 2, l, m, n] = dot(u[i, j, :, l, m, n], lr)
            u_face[i, j, 3, l, m, n] = dot(u[i, j, l, :, m, n], lr)
            u_face[i, j, 4, l, m, n] = dot(u[i, j, :, l, m, n], ll)
        end
    end

    @inbounds @threads for o = 1:2
        for n = 1:nq, m = 1:4, l = 1:nsp, j in axes(u_face, 2), i in axes(u_face, 1)
            f_face[i, j, 1, l, m, n, o] = dot(f[i, j, l, :, m, n, o], ll)
            f_face[i, j, 2, l, m, n, o] = dot(f[i, j, :, l, m, n, o], lr)
            f_face[i, j, 3, l, m, n, o] = dot(f[i, j, l, :, m, n, o], lr)
            f_face[i, j, 4, l, m, n, o] = dot(f[i, j, :, l, m, n, o], ll)
        end
    end

    @inbounds @threads for k = 1:nsp
        for j = 1:ny, i = 1:nx+1
            fw = zeros(4, nq)
            for l = 1:nq
                tmp = @view fw[:, l]
                uL = @view u_face[i-1, j, 2, k, :, l]
                uR = @view u_face[i, j, 4, k, :, l]
                flux_hll!(tmp, uL, uR, γ, 1.0)
            end
            fx_interaction[i, j, k, :, :] .= fw
        end
    end
    @inbounds @threads for k = 1:nsp
        for j = 1:ny+1, i = 1:nx
            fw = zeros(4, nq)
            for l = 1:nq
                tmp = @view fw[:, l]
                uL = local_frame(u_face[i, j-1, 3, k, :, l], 0.0, 1.0)
                uR = local_frame(u_face[i, j, 1, k, :, l], 0.0, 1.0)
                flux_hll!(tmp, uL, uR, γ, 1.0)
                tmp .= global_frame(tmp, 0.0, 1.0)
            end
            fy_interaction[i, j, k, :, :] .= fw
        end
    end

    @inbounds @threads for n = 1:nq
        for m = 1:4, l = 1:nsp, k = 1:nsp, j = 1:ny, i = 1:nx
            rhs1[i, j, k, l, m, n] = dot(f[i, j, :, l, m, n, 1], lpdm[k, :])
        end
    end
    @inbounds @threads for n = 1:nq
        for m = 1:4, l = 1:nsp, k = 1:nsp, j = 1:ny, i = 1:nx
            rhs2[i, j, k, l, m, n] = dot(f[i, j, k, :, m, n, 2], lpdm[l, :])
        end
    end

    @inbounds @threads for n = 1:nq
        for m = 1:4, l = 1:nsp, k = 1:nsp, j = 1:ny, i = 1:nx
            du[i, j, k, l, m, n] =
                -(
                    rhs1[i, j, k, l, m, n] + rhs2[i, j, k, l, m, n] +
                    (fx_interaction[i, j, l, m, n] * iJ[i, j][k, l][1, 1] - f_face[i, j, 4, l, m, n, 1]) * dhl[k] +
                    (fx_interaction[i+1, j, l, m, n] * iJ[i, j][k, l][1, 1] - f_face[i, j, 2, l, m, n, 1]) * dhr[k] +
                    (fy_interaction[i, j, k, m, n] * iJ[i, j][k, l][2, 2] - f_face[i, j, 1, k, m, n, 2]) * dhl[l] +
                    (fy_interaction[i, j+1, k, m, n] * iJ[i, j][k, l][2, 2] - f_face[i, j, 3, k, m, n, 2]) * dhr[l]
                )
        end
    end

    return nothing
end

begin
    f = OffsetArray{Float64}(undef, 0:ks.ps.nx+1, 0:ks.ps.ny+1, ks.ps.deg+1, ks.ps.deg+1, 4, uq.nq, 2)
    u_face = OffsetArray{Float64}(undef, 0:ks.ps.nx+1, 0:ks.ps.ny+1, 4, ks.ps.deg+1, 4, uq.nq)
    f_face = OffsetArray{Float64}(undef, 0:ks.ps.nx+1, 0:ks.ps.ny+1, 4, ks.ps.deg+1, 4, uq.nq, 2)
    fx_interaction = zeros(ks.ps.nx+1, ks.ps.ny, ks.ps.deg+1, 4, uq.nq)
    fy_interaction = zeros(ks.ps.nx, ks.ps.ny+1, ks.ps.deg+1, 4, uq.nq)
    rhs1 = zeros(ks.ps.nx, ks.ps.ny, ks.ps.deg+1, ks.ps.deg+1, 4, uq.nq)
    rhs2 = zeros(ks.ps.nx, ks.ps.ny, ks.ps.deg+1, ks.ps.deg+1, 4, uq.nq)
end

p = (f, u_face, f_face, fx_interaction, fy_interaction, rhs1, rhs2,
    ps.iJ, ps.ll, ps.lr, ps.dhl, ps.dhr, ps.dl, ks.gas.γ, uq)
tspan = (0.0, 1.0)
dt = 0.001
nt = tspan[2] ÷ dt |> Int

# initial condition
u0 = OffsetArray{Float64}(undef, 0:ps.nx+1, 0:ps.ny+1, ps.deg+1, ps.deg+1, 4, uq.nq)
begin
    gam = gas.γ
    MaL = gas.Ma
    MaR = sqrt((MaL^2 * (gam - 1.0) + 2.0) / (2.0 * gam * MaL^2 - (gam - 1.0)))
    ratioT =
        (1.0 + (gam - 1.0) / 2.0 * MaL^2) * (2.0 * gam / (gam - 1.0) * MaL^2 - 1.0) /
        (MaL^2 * (2.0 * gam / (gam - 1.0) + (gam - 1.0) / 2.0))
    t1 = [1.0, MaL * sqrt(gam / 2.0), 0.0, 1.0]
    t2 = [
        t1[1] * (gam + 1.0) * MaL^2 / ((gam - 1.0) * MaL^2 + 2.0),
        MaR * sqrt(gam / 2.0) * sqrt(ratioT),
        0.0,
        t1[end] / ratioT,
    ]

    for i in axes(u0, 1), j in axes(u0, 2), k in axes(u0, 3), l in axes(u0, 4)
        primRan = zeros(4, uq.nq)
        Ma = uq.pceSample .* ks.gas.Ma

        for n = 1:uq.nq
            if ps.x[i, j] <= ps.x1 * 0.25
                primRan[:, n] .= [t2[1], t1[2] - t2[2], 0.0, t2[end]]
            else
                primRan[:, n] .= [t1[1], 0.0, 0.0, t1[end]]

                tmp = @view primRan[:, n]
                vortex_ic!(tmp, ks.gas.γ, ps.xpg[i, j, k, l, 1], ps.xpg[i, j, k, l, 2])
            end
        end

        uRan = uq_prim_conserve(primRan, ks.gas.γ, uq)
        for m = 1:4
            u0[i, j, k, l, m, :] .= uRan[m, :]
        end
    end
end

prob = ODEProblem(dudt!, u0, tspan, p)
itg = init(prob, Midpoint(), save_everystep = false, adaptive = false, dt = dt)

@showprogress for iter = 1:nt
    step!(itg)

    # filter
    @inbounds @threads for j in axes(itg.u, 2)
        for i in axes(itg.u, 1)
            #ũ = ps.iV * reshape(itg.u[i, j, 1:3, 1:3, 1, :], 9, :)
            #su = maximum([sum(ũ[3:end, j].^2) / sum(ũ[:, j].^2) for j = 1:uq.nq])
            #isShock = shock_detector(log10(su), ps.deg)

            if true#isShock
                #λ = sqrt(su) * 2 #2e-5
                λ1 = 5e-5
                #λ2 = 1e-12#5
                for s = 1:4, ss = 1:uq.nq
                    û = ps.iV * reshape(itg.u[i, j, 1:3, 1:3, s, ss], 9)
                    FR.modal_filter!(û, λ1; filter = :l2)
                    itg.u[i, j, :, :, s, ss] .= reshape(ps.V * û, 3, 3)
                end
            end
        end
    end

    # boundary
    itg.u[:, 0, :, :, :, :] .= itg.u[:, 1, :, :, :, :]
    itg.u[:, ps.ny+1, :, :, :, :] .= itg.u[:, ps.ny, :, :, :, :]
    itg.u[ps.nx+1, :, :, :, :, :] .= itg.u[ps.nx, :, :, :, :, :]

    if iter % 100 == 0
        t = round(itg.t, digits=3)
        filename = "iter_" * string(t) * ".jld2"
        u = itg.u
        @save filename u
    end
end

begin
    x = zeros(ps.nx*(ps.deg+1), ps.ny*(ps.deg+1))
    y = zeros(ps.nx*(ps.deg+1), ps.ny*(ps.deg+1))
    sol = zeros(ps.nx*(ps.deg+1), ps.ny*(ps.deg+1), 4, 2)

    for i = 1:ps.nx, j = 1:ps.ny
        idx0 = (i - 1) * (ps.deg+1)
        idy0 = (j - 1) * (ps.deg+1)

        for k = 1:ps.deg+1, l = 1:ps.deg+1
            idx = idx0 + k
            idy = idy0 + l
            x[idx, idy] = ps.xpg[i, j, k, l, 1]
            y[idx, idy] = ps.xpg[i, j, k, l, 2]

            primRan = uq_conserve_prim(itg.u[i, j, k, l, :, :], ks.gas.γ, uq)
            primRan[4, :] .= 1 ./ primRan[4, :]
            primChaos = zeros(4, uq.nm+1)
            for ii = 1:4
                primChaos[ii, :] .= ran_chaos(primRan[ii, :], uq)
            end

            for s = 1:4
                sol[idx, idy, s, 1] = mean(primChaos[s, :], uq.op)
                sol[idx, idy, s, 2] = std(primChaos[s, :], uq.op)
            end
        end
    end

    #contourf(x, y, sol[:, :, 1], aspect_ratio=1, legend=true)
    #plot!(x[:, 1], sol[:, end÷2+1, 1])
end

#contourf(x[:, 1], y[1, :], sol[:, :, 4, 1]', aspect_ratio=1, legend=true)
#plot(x[:, 1], sol[:, end÷2+1, 1, 1])

u = itg.u
@save "sol.jld2" x sol u

max_stars_repo_name: CSMMLab/FRSG

max_issues_repo_name: CSMMLab/FRSG

max_forks_repo_name: CSMMLab/FRSG


text:
struct PendulumParams{T}
    maxspeed::T
    maxtorque::T
    g::T
    m::T
    l::T
    dt::T
    maxsteps::Int64
end

mutable struct Pendulum{T} <: AbstractEnv
    params::PendulumParams{T}
    actionspace::BoxSpace{T}
    observationspace::BoxSpace{T}
    state::Array{T, 1}
    done::Bool
    t::Int64
end

function Pendulum(; T = Float64, maxspeed = T(8), maxtorque = T(2),
                    g = T(10), m = T(1), l = T(1), dt = T(.05), maxsteps = 200)
    high = T.([1, 1, maxspeed])
    env = Pendulum(PendulumParams(maxspeed, maxtorque, g, m, l, dt, maxsteps),
                   BoxSpace(-2., 2.),
                   BoxSpace(-high, high),
                   zeros(T, 2), false, 0)
    reset!(env)
    env
end

actionspace(env::Pendulum) = env.actionspace

pendulumobservation(s) = [cos(s[1]), sin(s[1]), s[2]]
anglenormalize(x) = ((x + pi) % (2*pi)) - pi

getstate(env::Pendulum) = (observation=pendulumobservation(env.state), isdone=env.done)

function reset!(env::Pendulum{T}) where T
    env.state[:] = 2 * rand(T, 2) .- 1
    env.t = 0
    env.done = false
    (observation=pendulumobservation(env.state),)
end

interact!(env::Pendulum, a::AbstractArray) = interact!(env, a[1])
function interact!(env::Pendulum, a)
    env.t += 1
    th, thdot = env.state
    a = clamp(a, -env.params.maxtorque, env.params.maxtorque)
    costs = anglenormalize(th)^2 + .1 * thdot^2 + .001 * a^2
    newthdot = thdot + (-3 * env.params.g/(2*env.params.l) * sin(th + pi) +
                        3 * a/(env.params.m * env.params.l^2)) * env.params.dt
    th += newthdot * env.params.dt
    newthdot = clamp(newthdot, -env.params.maxspeed, env.params.maxspeed)
    env.state[1] = th
    env.state[2] = newthdot
    env.done = env.t >= env.params.maxsteps
    (observation=pendulumobservation(env.state), reward=-costs, isdone=env.done)
end


text:
using LinearAlgebra
using Random
using NNlib: softplus
using Roots # for inverse

################################################################################
#                            Planar and Radial Flows                           #
#             Ref: Variational Inference with Normalizing Flows,               #
#               D. Rezende, S. Mohamed(2015) arXiv:1505.05770                  #
################################################################################

###############
# PlanarLayer #
###############

# TODO: add docstring

mutable struct PlanarLayer{T1<:AbstractVector{<:Real}, T2<:Real} <: Bijector{1}
    w::T1
    u::T1
    b::T2
end

function get_u_hat(u, w)
    # To preserve invertibility
    x = w' * u
    return u .+ (planar_flow_m(x) - x) .* w ./ sum(abs2, w)   # from A.1
end

function PlanarLayer(dims::Int, container=Array)
    w = container(randn(dims))
    u = container(randn(dims))
    b = randn()
    return PlanarLayer(w, u, b)
end

planar_flow_m(x) = -1 + softplus(x)   # for planar flow from A.1
ψ(z, w, b) = (1 .- tanh.(w' * z .+ b).^2) .* w    # for planar flow from eq(11)

# An internal version of transform that returns intermediate variables
function _transform(flow::PlanarLayer, z::AbstractVecOrMat)
    _planar_transform(flow.u, flow.w, flow.b, z)
end
function _planar_transform(u, w, b, z)
    u_hat = get_u_hat(u, w)
    transformed = z .+ u_hat .* tanh.(w' * z .+ b) # from eq(10)
    return (transformed = transformed, u_hat = u_hat)
end

(b::PlanarLayer)(z) = _transform(b, z).transformed

function forward(flow::PlanarLayer, z::AbstractVecOrMat)
    transformed, u_hat = _transform(flow, z)
    # Compute log_det_jacobian
    psi = ψ(z, flow.w, flow.b) .+ zero(eltype(u_hat))
    if psi isa AbstractVector
        T = eltype(psi)
    else
        T = typeof(vec(psi))
    end
    log_det_jacobian::T = log.(abs.(1.0 .+ psi' * u_hat)) # from eq(12)
    return (rv = transformed, logabsdetjac = log_det_jacobian)
end

function (ib::Inverse{<: PlanarLayer})(y::AbstractVector{<:Real})
    flow = ib.orig
    u_hat = get_u_hat(flow.u, flow.w)
    T = promote_type(eltype(flow.u), eltype(flow.w), eltype(flow.b), eltype(y))
    TV = vectorof(T)
    # Define the objective functional; implemented with reference from A.1
    f(y) = alpha -> (flow.w' * y) - alpha - (flow.w' * u_hat) * tanh(alpha + flow.b)
    # Run solver
    alpha::T = find_zero(f(y), zero(T), Order16())
    z_para::TV = (flow.w ./ norm(flow.w, 2)) .* alpha
    return (y .- u_hat .* tanh.(flow.w' * z_para .+ flow.b))::TV
end
function (ib::Inverse{<: PlanarLayer})(y::AbstractMatrix{<:Real})
    flow = ib.orig
    u_hat = get_u_hat(flow.u, flow.w)
    T = promote_type(eltype(flow.u), eltype(flow.w), eltype(flow.b), eltype(y))
    TM = matrixof(T)
    # Define the objective functional; implemented with reference from A.1
    f(y) = alpha -> (flow.w' * y) - alpha - (flow.w' * u_hat) * tanh(alpha + flow.b)
    # Run solver
    alpha = mapvcat(eachcol(y)) do c
        find_zero(f(c), zero(T), Order16())
    end
    z_para::TM = (flow.w ./ norm(flow.w, 2)) .* alpha'
    return (y .- u_hat .* tanh.(flow.w' * z_para .+ flow.b))::TM
end

function matrixof(::Type{Vector{T}}) where {T <: Real}
    return Matrix{T}
end
function matrixof(::Type{T}) where {T <: Real}
    return Matrix{T}
end
function vectorof(::Type{Matrix{T}}) where {T <: Real}
    return Vector{T}
end
function vectorof(::Type{T}) where {T <: Real}
    return Vector{T}
end

logabsdetjac(flow::PlanarLayer, x) = forward(flow, x).logabsdetjac
isclosedform(b::Inverse{<:PlanarLayer}) = false


text:
import Test: @test, @testset

@testset "Radial propagation" begin
import Luna
import Luna: Grid, Maths, PhysData, Nonlinear, Ionisation, NonlinearRHS, Output, Stats, LinearOps, Plotting
import Luna.PhysData: wlfreq
import FFTW
import Hankel
import LinearAlgebra: norm

gas = :Ar
pres = 1.2
τ = 20e-15
λ0 = 800e-9
w0 = 40e-6
energy = 1e-12
L = 0.6
R = 4e-3
N = 128

grid = Grid.RealGrid(L, 800e-9, (400e-9, 2000e-9), 0.2e-12)
q = Hankel.QDHT(R, N, dim=2)

energyfun, energyfun_ω = Fields.energyfuncs(grid, q)

function prop(E, z)
    Eω = FFTW.rfft(E, 1)
    Eωk = q * Eω
    kzsq = @. (grid.ω/PhysData.c)^2 - (q.k^2)'
    kzsq[kzsq .< 0] .= 0
    kz = sqrt.(kzsq)
    @. Eωk *= exp(-1im * z * (kz - grid.ω/PhysData.c))
    Eω = q \ Eωk
    E = FFTW.irfft(Eω, length(grid.t), 1)
    return E
end

dens0 = PhysData.density(gas, pres)
densityfun(z) = dens0
ionpot = PhysData.ionisation_potential(gas)
ionrate = Ionisation.ionrate_fun!_PPTcached(gas, λ0)
responses = (Nonlinear.Kerr_field(PhysData.γ3_gas(gas)),
             Nonlinear.PlasmaCumtrapz(grid.to, grid.to, ionrate, ionpot))
linop = LinearOps.make_const_linop(grid, q, PhysData.ref_index_fun(gas, pres))
normfun = NonlinearRHS.const_norm_radial(grid, q, PhysData.ref_index_fun(gas, pres))

inputs = Fields.GaussGaussField(λ0=λ0, τfwhm=τ, energy=energy, w0=w0, propz=-0.3)

Eω, transform, FT = Luna.setup(grid, q, densityfun, normfun, responses, inputs)
output = Output.MemoryOutput(0, grid.zmax, 201)
Luna.run(Eω, grid, linop, transform, FT, output)

Erout = (q \ output.data["Eω"])
Iωr = abs2.(Erout)

ω0idx = argmin(abs.(grid.ω .- 2π*PhysData.c/λ0))

Iλ0 = Iωr[ω0idx, :, :]
λ0 = 2π*PhysData.c/grid.ω[ω0idx]
w1 = w0*sqrt(1+(L/2*λ0/(π*w0^2))^2)
Iλ0_analytic = Maths.gauss.(q.r, w1/2)*(w0/w1)^2 # analytical solution (in paraxial approx)
Ir = Maths.normbymax(Iλ0[:, end])
Ira = Maths.normbymax(Iλ0_analytic)
@test maximum(abs.(Ir .- Ira)/norm(Ir)) < 1.5e-4
end


text:
using ACE, LinearAlgebra, StaticArrays
using ACE: evaluate
using ACEtb
using Random: shuffle

# this should possibly go into the main ACEtb module, or possibly into
# the hamiltonian assembly code
function eval_bond(B, Rs, Zs, z0)
   r̂ = Rs[1] / norm(Rs[1])
   o = Rs[1]/2
   Rr = map( r_ -> (r = r_ - o; o + r - 2*dot(r,r̂)*r̂), Rs )
   Rr[1] = Rs[1]
   return evaluate(B, Rs, Zs, z0) + evaluate(B, Rr, Zs, z0)
end

# utility function for generating random isometries
function randsym(Rs)
   Rs1 = [ [Rs[1]]; shuffle(Rs[2:end]) ]
   K = randn(3, 3)
   K = K - K'
   Q = SMatrix{3,3}(rand([-1,1]) * exp(K)...)
   return [ Q * R for R in Rs1 ]
end

#---

# to create a basis with cylindrical symmetry start with a
# standard ACE basis
r0 = rnn(:Al)
B = ACE.Utils.ace_basis( species = [:X, :Al], N = 1,
                         pin = 0, pcut = 0)

# define the cylindrical cut-off function
# Fcut = ACEtb.Bonds.BondCutoff(r0 = rnn(:Al))
Fcut = ACEtb.Bonds.BondCutoff(r0 = nothing, rcut = 2.5*r0, renv = 1.9*r0, zenv=1.9*r0, pcut = 2)

# convert the radial ACE basis into a cylindrical ACE basis.
Bbonds = ACEtb.Bonds.RPIBonds(B, Fcut)

#---

#---

# Testing symmetry
for ntest = 1:10
   Rs, Zs, z0 = ACE.Random.rand_nhd(10, B.pibasis.basis1p.J, :Al)
   Zs[1] = 0
   b0 = eval_bond(Bbonds, Rs, Zs, z0)
   Rs1 = randsym(Rs)
   b1 = eval_bond(Bbonds, Rs1, Zs, z0)
   @show b1 ≈ b0
end

#---
# Plotting bond basis values
xgr = range(-2*r0, 3*r0, length = 40)
ygr = range(-2*r0, 2*r0, length = 40)
R0 = SVector([r0, 0.0, 0.0]...)
Zs = AtomicNumber.([:X, :Al])
z0 = AtomicNumber(:X)
X = []
Y = []
V = []

for x in xgr, y in ygr
   R1 = SVector([x, y, 0.0]...)
   # if norm(R1) < 0.5*r0 || norm(R1 - R0) < 0.5*r0
   #    continue
   # end
   Rs = [R0, R1]
   b0 = eval_bond(Bbonds, Rs, Zs, z0)
   # b0 = evaluate(Bbonds, Rs, Zs, z0)
   push!(V,b0)
   push!(X,x)
   push!(Y,y)
end

vplot = [V[k][9] for k in 1:length(V)]
#using Plots
#scatter(X,Y, marker_z = vplot, color = :jet, clims = (-0.7, 0.7))
#plot!([0.0, r0], [0.0, 0.0], lw=4, m=:o, ms=6, c=:red, label = "")
# #---
# minimum(abs.(vplot))
# maximum(abs.(vplot))
@show [extrema([V[k][i] for k in 1:length(V)]) for i = 1:32]
#
# v[1]
# display([ v[1] v[2] ][1:16,:])


text:
using QuantumOptics: basisstate

export VibrationalMode

"""
    VibrationalMode(
            ν::Real, mode_structure::Vector{Real}, δν::Union{Function,Real}=0.; N::Int=10,
            axis::NamedTuple{(:x,:y,:z)}=ẑ
        )

**user-defined fields**
* `ν::Real`: frequency [Hz]
* `mode_structure::Vector{Real}`: The normalized eigenvector describing the collective motion
        of the ions belonging to this mode.
* `δν::Union{Function,Real}`: Either a function describing time-dependent fluctuations of `ν`
        or a real number which will be converted to the constant function `t -> δν`.
* `N::Int`: Highest level included in the Hilbert space.
* `axis::NamedTuple{(:x,:y,:z)}`: The axis of symmetry for the vibration. This must lie along
        one of the basis vectors `x̂`, `ŷ` or `ẑ`.
**derived fields**
* `shape::Vector{Int}`: Indicates dimension of used Hilbert space (`=[N+1]`).
* `_cnst_δν::Bool`: A Boolean flag signifying whether or not `δν` is a constant function.

Note: the iᵗʰ Fock state (|i⟩) can be obtained by indexing as `v=VibrationalMode(...); v[i]`
"""
mutable struct VibrationalMode <: IonSimBasis
    ν::Real
    mode_structure::Vector{Real}
    δν::Function
    N::Int
    shape::Vector{Int}
    axis::NamedTuple{(:x,:y,:z)}
    _cnst_δν::Bool
    function VibrationalMode(ν, mode_structure; δν::Tδν=0., N=10, axis=ẑ) where {Tδν}
        if Tδν <: Number
            _cnst_δν = true
            δνt(t) = δν
        else
            _cnst_δν = false
            δνt = δν
        end
        new(ν, mode_structure, δνt, N, [N+1], axis, _cnst_δν)
    end
end

function Base.:(==)(b1::T, b2::T) where {T<:VibrationalMode}
    (
        b1.ν == b2.ν &&
        b1.mode_structure == b2.mode_structure &&
        b1.N == b2.N &&
        b1.shape == b2.shape &&
        b1.axis == b2.axis
    )
end

# suppress long output
Base.show(io::IO, V::VibrationalMode) = print(io,
    "VibrationalMode(ν=$(round(V.ν,sigdigits=4)), axis=$(_print_axis(V.axis)), N=$(V.N))")

function Base.setproperty!(V::VibrationalMode, s::Symbol, v::Tv) where {Tv}
    if s == :mode_structure || s == :axis || s == :_cnst_δν
        return
    end
    if s == :N
        @assert Tv <: Int "N must be a positive integer"
        @assert v >= 0 "N must be a nonnegative integer"
        Core.setproperty!(V, :N, v)
        Core.setproperty!(V, :shape, Int[v+1])
    elseif s == :shape
        return
    elseif s == :δν
        if Tv <: Number
            _cnst_δν = true
            vt(t) = v
        else
            _cnst_δν = false
            vt = v
        end
        Core.setproperty!(V, s, vt)
        Core.setproperty!(V, :_cnst_δν, _cnst_δν)
        return
    end
    Core.setproperty!(V, s, v)
end

function Base.getindex(V::VibrationalMode, n::Int)
    @assert 0 <= n <= V.N "n ∉ [0, $(V.N+1)]"
    basisstate(V, n+1)
end


text:
struct NeRFRenderer
    coarse::NeRFModel
    fine::NeRFModel
end

Flux.@functor NeRFRenderer

function render_rays(nr::NeRFRenderer, batch; coarse_ts, fine_ts, bbox_min, bbox_max, background, rng::AbstractRNG=Random.GLOBAL_RNG)
    t_min, t_max, mask = batched_t_range(batch; bbox_min, bbox_max)
    coarse_ts = stratified_sampling(t_min, t_max, mask, coarse_ts; rng)
    all_points = points(coarse_ts, batch)
    direction_batch = _repeat(batch[:, 2:2, :], 1, size(all_points, 2), 1)
    coarse_densities, coarse_rgbs = nr.coarse(reshape(all_points, size(all_points, 1), :), reshape(direction_batch, size(direction_batch, 1), :))
    coarse_rgbs = reshape(coarse_rgbs, size(all_points)...)
    coarse_densities = reshape(coarse_densities, size(all_points)[2:end]...)
    coarse_outputs = render_rays(coarse_ts, coarse_densities, coarse_rgbs, background)

    fine_ts = fine_sampling(coarse_ts, fine_ts, coarse_densities; rng)
    all_points = points(fine_ts, batch)
    direction_batch = _repeat(batch[:, 2:2, :], 1, size(all_points, 2), 1)
    fine_densities, fine_rgbs = nr.fine(reshape(all_points, size(all_points, 1), :), reshape(direction_batch, size(direction_batch, 1), :))
    fine_rgbs = reshape(fine_rgbs, size(all_points)...)
    fine_densities = reshape(fine_densities, size(all_points)[2:end]...)
    fine_outputs = render_rays(fine_ts, fine_densities, fine_rgbs, background)
    return (coarse=coarse_outputs, fine=fine_outputs)
end

function batched_t_range(batch; bbox_min, bbox_max, eps=Float32(1e-8), min_t_range=Float32(1e-3))
    bbox = cat(bbox_min, bbox_max; dims=ndims(bbox_min)+1)

    # batch: 3x2xB
    origin = batch[:, 1:1, :]
    direction = batch[:, 2:2, :]

    offsets = bbox .- origin # 3 x B
    @check ndims(offsets) == 3

    ts = offsets ./ (direction .+ eps)
    # search sorted
    ts = hcat(minimum(ts, dims=2), maximum(ts, dims=2))

    @check size(ts) == size(offsets)
    # Find overlapping and apply constraints
    max_of_min = maximum(ts[:, 1, :], dims=1)
    min_t = relu.(max_of_min)
    max_t = minimum(ts[:, 2, :], dims=1)
    max_t_clipped = max.(max_t, min_t .+ min_t_range)
    real_range = vcat(min_t, max_t_clipped)
    null_range = Flux.Adapt.adapt(typeof(real_range), reshape(vcat(0, min_t_range), 2, 1))
    mask = min_t .< max_t
    bounds = ifelse.(mask, real_range, null_range)
    @check ndims(bounds) == 2
    return bounds[1, :], bounds[2, :], dropdims(mask, dims=1)
end

struct RaySamples
    t_min
    t_max
    mask
    ts
end

function stratified_sampling(t_min::AbstractVector, t_max::AbstractVector, mask, count::Int; rng=Random.GLOBAL_RNG)
    bin_size = (t_max .- t_min) ./ count
    bin_size = unsqueeze(bin_size, dims=1)
    tmp = range(0, count-1)
    bin_starts = unsqueeze(tmp, dims=2) .* bin_size .+ unsqueeze(t_min, dims=1)
    @check size(bin_starts) == (count, size(t_min, 1))
    randoms = Flux.Zygote.@ignore rand(rng, Float32, size(bin_starts)) .* bin_size
    ts = bin_starts .+ randoms
    return RaySamples(t_min, t_max, mask, ts)
end

"""
For each ray, compute the points at all timesteps

- `rays`: 3x2xB batch of rays

Returns a batch of points of shape 3xTxB
"""
function points(rs::RaySamples, rays)
    return rays[:, begin:begin, :] .+ (rays[:, begin+1:end, :] .* unsqueeze(rs.ts, dims=1))
end

function starts(rs::RaySamples)
    t_mid = (rs.ts[begin+1:end,:] .+ rs.ts[begin:end-1,:]) ./ 2
    return vcat(reshape(rs.t_min, 1, size(rs.t_min)...), t_mid)
end

function ends(rs::RaySamples)
    t_mid = (rs.ts[begin+1:end, :] .+ rs.ts[begin:end-1, :]) ./ 2
    return vcat(t_mid, reshape(rs.t_max, 1, size(rs.t_max)...))
end

function deltas(rs::RaySamples)
    ends(rs) .- starts(rs)
end

function termination_probs(rs::RaySamples, densities)
    density_dt = densities .* deltas(rs)

    # Integral of termination probabilities over time.
    acc_densities_cur = cumsum(density_dt, dims=1) # TODO check
    tmp1 = Flux.Adapt.adapt(typeof(densities), zeros(eltype(densities), 1, size(acc_densities_cur,2)))
    acc_densities_prev = vcat(tmp1, acc_densities_cur)
    prob_survive = exp.(-acc_densities_prev)

    # Probability of terminating at time t given we made it to time t
    tmp2 = ones(eltype(densities), 1, size(rs.ts)[end])
    prob_terminate = vcat((1 .- exp.(-density_dt)), tmp2)
    return prob_survive .* prob_terminate
end

"""

Volumetric rendering given density and color samples along a batch of rays.
"""
function render_rays(rs::RaySamples, densities::AbstractMatrix{T}, rgbs::AbstractArray{T,3}, background::AbstractVector{T}) where {T<:AbstractFloat}
    probs = termination_probs(rs, densities)
    colors = hcat(rgbs, _repeat(reshape(background, size(background, 1), 1, 1) , 1, 1, size(rgbs)[end])) # TODO
    fg = dropdims(sum(unsqueeze(probs, dims=1) .* colors; dims=2), dims=2)
    bg = unsqueeze(background, dims=2)
    ms = unsqueeze(rs.mask, dims=1)
    return @. ifelse(ms, fg, bg)
end

function fine_sampling(rs::RaySamples, count::Int, densities; combine::Bool=false, eps=Float32(1e-8), rng::AbstractRNG=Random.GLOBAL_RNG)
    w = termination_probs(rs, densities)[begin:end-1, :] .+ eps

    # Setup an inverse CDF for inverse transform sampling
    xs = cumsum(w; dims=1)
    tmp = Flux.Adapt.adapt(typeof(densities), zeros(eltype(densities), 1, size(rs.ts)[end]))
    xs =  vcat(tmp, xs)
    xs = xs ./ xs[end:end, :]
    ys = vcat(reshape(rs.t_min, 1, size(rs.t_min)...), ends(rs))

    # Evaluate the inverse CDF at quasi-random points.
    input_samples = stratified_sampling(rs.t_min, rs.t_max, rs.mask, count; rng=rng)
    new_ts = batched_interpolate(input_samples.ts, xs, ys)

    if combine
        # TODO: sort along a dimension is not differentiable
        error("not-implemented")
        combined = vcat(rs.ts, new_ts)
        new_ts = sort(combined, dims=1)
    end
    return RaySamples(rs.t_min, rs.t_max, rs.mask, new_ts)
end

# function _sort(xs::AbstractMatrix)
#     xsb = unbatch(xs)
#     sx = sort.(xsb)
#     stack(sx, dims=2)
# end


text:
const JLL_ENV_GCC_VERSIONS = VersionNumber[
    v"4.8.5",
    v"5.2.0",
    v"6.1.0",
    v"7.1.0",
    v"8.1.0",
    v"9.1.0",
    v"10.2.0",
    v"11.0.0-iains",
]

const GCC_MIN_VER = minimum(JLL_ENV_GCC_VERSIONS)


text:
struct SymWoodbury{T,AType,BType,DType,DpType} <: AbstractWoodbury{T}
    A::AType
    B::BType
    D::DType
    Dp::DpType
    tmpN1::Vector{T}
    tmpN2::Vector{T}
    tmpk1::Vector{T}
    tmpk2::Vector{T}
end

"""
    W = SymWoodbury(A, B, D)

Represent a matrix of the form `W = A + BDBᵀ`, where `A` and `D` are symmetric.
Equations `Wx = b` will be solved using the
[Woodbury matrix identity](https://en.wikipedia.org/wiki/Woodbury_matrix_identity).

If your main goal is to solve equations, it can be advantageous to supply
`A` as a factorization (e.g., `SymWoodbury(cholesky(A), B, D)` if `A` is a positive
semidefinite matrix).

This checks that `A` is symmetric; you can elide the check by passing a `Symmetric` matrix
or factorization.

See also [Woodbury](@ref).
"""
function SymWoodbury(A, B::AbstractVecOrMat, D)
    @noinline throwdmm(B, D, A) = throw(DimensionMismatch("Sizes of B ($(size(B))) and/or D ($(size(D))) are inconsistent with A ($(size(A)))"))

    n = size(A, 1)
    k = size(B, 2)
    if size(A, 2) != n || size(B, 1) != n || size(D,1) != k || size(D,2) != k
        throwdmm(B, D, A)
    end
    if !isa(A, Cholesky) && !isa(A, BunchKaufman) && !isa(A, LDLt)
        issymmetric(A) || throw(ArgumentError("A must be symmetric"))
    end
    issymmetric(D) || throw(ArgumentError("D must be symmetric"))
    Dp = safeinv(safeinv(D) .+ B'*(A\B))
    # temporary space for allocation-free solver (vector RHS only)
    T = typeof(float(zero(eltype(A)) * zero(eltype(B)) * zero(eltype(D))))
    tmpN1 = Vector{T}(undef, n)
    tmpN2 = Vector{T}(undef, n)
    tmpk1 = Vector{T}(undef, k)
    tmpk2 = Vector{T}(undef, k)

    SymWoodbury(A, B, D, Dp, tmpN1, tmpN2, tmpk1, tmpk2)
end

convert(::Type{W}, O::SymWoodbury) where {W<:Woodbury} = Woodbury(O.A, O.B, O.D, O.B')

convert(::Type{Matrix{T}}, W::SymWoodbury) where {T} = convert(Matrix{T}, symmetrize!(Matrix(W.A) + W.B*W.D*W.B'))

# This allows a common interface with Woodbury
function Base.getproperty(O::SymWoodbury, name::Symbol)
    # SymWoodbury names
    name === :A && return getfield(O, :A)
    name === :B && return getfield(O, :B)
    name === :D && return getfield(O, :D)
    name === :Dp && return getfield(O, :Dp)
    name === :tmpN1 && return getfield(O, :tmpN1)
    name === :tmpN2 && return getfield(O, :tmpN2)
    name === :tmpk1 && return getfield(O, :tmpk1)
    name === :tmpk2 && return getfield(O, :tmpk2)
    # Woodbury names
    name === :U && return getfield(O, :B)
    name === :C && return getfield(O, :D)
    name === :Cp && return getfield(O, :Dp)
    name === :V && return getfield(O, :B)'
    error("property ", name, " not defined")
end

inv_invD_BtX(invD, B, X) = safeinv(invD - B'*X)
inv_invD_BtX(invD, B::AbstractVector, X) = safeinv(invD - dot(B,X))

function calc_inv(A, B, D)
    W = safeinv(A)
    X = W*B
    Z = inv_invD_BtX(-safeinv(D), B, X)
    SymWoodbury(symmetrize!(W), X, symmetrize!(Z))
end

inv(O::SymWoodbury{T,AType,BType,DType}) where {T<:Any, AType<:Any, BType<:AbstractVector, DType<:Real} =
  calc_inv(O.A, O.B, O.D)

inv(O::SymWoodbury{T,AType,BType,DType}) where {T<:Any, AType<:Any, BType<:Any, DType<:AbstractMatrix} =
  calc_inv(O.A, O.B, O.D)

# D is typically small, so this is acceptable.
inv(O::SymWoodbury{T,AType,BType,DType}) where {T<:Any, AType<:Any, BType<:Any, DType<:SparseMatrixCSC} =
  calc_inv(O.A, O.B, Matrix(O.D))

_ldiv!(dest, W::SymWoodbury, A::AbstractMatrix, B) = _ldiv!(dest, W, defaultfactor(A), B)

defaultfactor(A) = bunchkaufman(A, true)
defaultfactor(A::SymTridiagonal) = ldlt(A)
defaultfactor(A::SparseMatrixCSC) = lu(A)

"""
    partialInv(O)

Get the factors `(X, Z)` in `Y + XZXᵀ == inv(A + BDBᵀ)`
"""
function partialInv(O::SymWoodbury)
    X = (O.A)\O.B
    invD = -1*inv(O.D)
    Z = inv_invD_BtX(invD, O.B, X)
    return (X,Z)
end

+(O::SymWoodbury, M::SymWoodbury)    = SymWoodbury(O.A + M.A, [O.B M.B],
                                                   cat(O.D,M.D; dims=(1,2)) )
*(α::Real, O::SymWoodbury)           = SymWoodbury(α*O.A, O.B, α*O.D)
+(O::SymWoodbury, M::AbstractMatrix) = SymWoodbury(O.A + M, O.B, O.D)

Base.copy(O::SymWoodbury{T}) where {T} = SymWoodbury(copy(O.A), copy(O.B), copy(O.D))

function square(O::SymWoodbury)
    A  = O.A^2
    AB = O.A*O.B
    Z  = [(AB + O.B) (AB - O.B)]
    R  = symmetrize!(O.D*(O.B'*O.B)*O.D/4)
    D  = [ O.D/2 + R  -R
          -R          -O.D/2 + R ]
    SymWoodbury(A, Z, D)
end

"""
The product of two SymWoodbury matrices will generally be a Woodbury Matrix,
except when they are the same, i.e. the user writes A'A or A*A' or A*A.

Z(A + B*D*Bᵀ) = ZA + ZB*D*Bᵀ

This package will not support support left multiplication by a generic
matrix, to keep return types consistent.
"""
function *(O1::SymWoodbury, O2::SymWoodbury)
    if (O1 === O2)
        return square(O1)
    elseif O1.A == O2.A && O1.B == O2.B && O1.D == O2.D
        return square(O1)
    else
        throw(MethodError("ERROR: To multiply two non-identical SymWoodbury matrices, first convert to Woodbury."))
    end
end

conjm(O::SymWoodbury, M) = SymWoodbury(symmetrize!(M*O.A*M'), M*O.B, symmetrize!(O.D))

Base.getindex(O::SymWoodbury, I::UnitRange, I2::UnitRange) =
  SymWoodbury(O.A[I,I], O.B[I,:], O.D);

adjoint(O::SymWoodbury{T}) where T<:Real = O   # it's Hermitian
transpose(O::SymWoodbury) = O

function show(io::IO, W::SymWoodbury)
    println(io, "Symmetric Woodbury factorization:\nA:")
    show(io, MIME("text/plain"), W.A)
    print(io, "\nB:\n")
    Base.print_matrix(IOContext(io,:compact=>true), W.B)
    print(io, "\nD: ", W.D)
end

function symmetrize!(A::AbstractMatrix)
    axs = axes(A)
    axs[1] == axs[2] || error("matrix must be square")
    for j in axs[2], i in j+1:last(axs[1])
        a = (A[i,j] + A[j,i])/2
        A[i,j] = A[j,i] = a
    end
    return A
end

symmetrize!(x::Real) = x


text:
# TODO move this constructor to mcts.jl
function MCTS.Env(gspec::AbstractGameSpec, oracle, params::MctsParams; S=nothing)
	return MCTS.Env(gspec, oracle;
	gamma=params.gamma,
	cpuct=params.cpuct,
	noise_ϵ=params.dirichlet_noise_ϵ,
	noise_α=params.dirichlet_noise_α,
	prior_temperature=params.prior_temperature,
	S=S)
end

"""
		MuPlayer{P,D,R} <: AbstractPlayer

- `prediction_oracle` 		f(sᵏ) -> (pᵏ,vᵏ)
- `dynamics_oracle`				g(sᵏ⁻¹,aᵏ) -> (rᵏ,sᵏ)
- `representation_oracle` h(o) -> (s⁰)
- `mcts_params::MctsParams`
- `timeout::Union{Float64, Nothing}` time that MCTS has for thinking
		nothing means infinite (till num_iters_per_turn ends)
"""
# struct MuPlayer{P,D,R} <: AbstractPlayer
struct MuPlayer{I,R} <: AbstractPlayer
	# mcts :: M
	# prediction_oracle :: P
	# dynamics_oracle :: D
	# representation_oracle :: R
	initial_inference :: I
	recurrent_inference :: R
	mcts_params :: MctsParams
	timeout :: Union{Float64, Nothing}
	# niters :: Int
	# τ :: AbstractSchedule{Float64} # Temperature
	# function MuPlayer(mcts::MCTS.Env, representation_oracle, dynamics_oracle; τ, niters, timeout=nothing)
	function MuPlayer((init, recur), mcts_params::MctsParams; timeout=nothing)
		@assert mcts_params.num_iters_per_turn > 0
		@assert isnothing(timeout) || timeout > 0
		new{typeof(init), typeof(recur)}(init, recur, mcts_params, timeout)
	end
end

# # constructor with MuNetwork
# function MuPlayer(
# 	μnetwork::Union{MuNetwork,NamedTuple}, params::MctsParams; timeout=nothing)
#   return MuPlayer(μnetwork.f, μnetwork.g, μnetwork.h, params;
# 	# niters=params.num_iters_per_turn,
# 	# τ=params.temperature,
# 	timeout=timeout)
# end

function AlphaZero.player_temperature(p::MuPlayer, game, turn)
	return p.mcts_params.temperature[turn]
end

# function reset_player!(player::Union{MctsPlayer,MuPlayer})
# 	MCTS.reset!(player.mcts)
# end

# normalize policy - probs of legal actions sum up to 1, and illegals are 0
function normalize_p(P, actions_mask)
	P = P .* actions_mask # Zygote doesn't work with .*=
	sp = sum(P, dims=1)
	P = P ./ (sp .+ eps(eltype(P)))
	return P
end

# ? does rootinfo.vest should be there, or just
# ? sum without γ
# compute_rootvalue(rootinfo::MCTS.StateInfo, γ) = rootinfo.Vest + γ*sum([st.W for st in rootinfo.stats])
# compute_rootvalue(ri::MCTS.StateInfo, γ) = max(st.W/st.N for st in ri.stats)
compute_rootvalue(ri::MCTS.StateInfo) = sum(st.W for st in ri.stats) / MCTS.Ntot(ri)
# TODO compare max and sum

"""
		think(::MuPlayer, game)

Return avaliable actions, policy, and rootvalue
"""
function AlphaZero.think(plr::MuPlayer, game)
	#initial inference:  h(o) → s⁰,  f(s⁰) → (p⁰,v⁰)
	# rootstate = p.representation_oracle(GI.current_state(game))
	# (P⁰, V⁰) = p.prediction_oracle(rootstate)
	p⁰, v⁰, s⁰ = plr.initial_inference(GI.current_state(game))
	actions_mask = GI.actions_mask(game)
	p⁰ = normalize_p(p⁰, actions_mask)[actions_mask]

	recur_res = Ref((p⁰, v⁰, s⁰, zero(v⁰))) # g and f are distincted in mcts.jl, so there is need to pass by reference (or some kind of channel)
	prediction_recurrent_oracle(state) = state==recur_res[][3] ? recur_res[][1:2] : @error "states don't match" state recur_res[][3]

	mcts = MCTS.Env(GI.spec(game), prediction_recurrent_oracle, plr.mcts_params, S=typeof(s⁰))
	mcts.tree[s⁰] = MCTS.init_state_info(p⁰,v⁰,mcts.prior_temperature)

	mugame = MuGameEnvWrapper(
		game,
		#	CachedOracle(oracle, Statetype, Actiontype, Rewardtype)
		CachedOracle(plr.recurrent_inference,typeof(recur_res[]),typeof(s⁰),GI.action_type(mcts.gspec)),
		s⁰,
		true, #isroot
		GI.white_playing(game),
		0f0,
		recur_res)
	niters = plr.mcts_params.num_iters_per_turn - 1 # -1, because (P⁰,V⁰) is already there
	if isnothing(plr.timeout) # Fixed number of MCTS simulations
		MCTS.explore!(mcts, mugame, niters)
	else # Run simulations until timeout
		start = time()
		while time() - start < plr.timeout
		MCTS.explore!(mcts, mugame, niters)
		end
	end
	rootvalue = compute_rootvalue(mcts.tree[s⁰])
	actions, π_target = MCTS.policy(mcts, mugame)
	return actions, π_target, rootvalue
end

"""
		play_game(gspec, ::MuPlayer)

overloaded AlphaZero.play_game() function,
that additionally push selected action, and rootvalue into trace
"""
function AlphaZero.play_game(gspec, player::Union{MuPlayer,MinMax.Player,MctsPlayer}; flip_probability=0., initstate=nothing)
  game = isnothing(initstate) ? GI.init(gspec) : GI.init(gspec, initstate)
	# @debug "initstate" initstate GI.current_state(game)
  trace = MuTrace(GI.current_state(game))
  while true
    if GI.game_terminated(game)
      return trace
    end
    if !iszero(flip_probability) && rand() < flip_probability
      GI.apply_random_symmetry!(game)
    end
    actions, π_target, rootvalue = think(player, game)
		rootvalue = GI.white_playing(game) ? rootvalue : -rootvalue# ? reversing rootvalue wp?
    τ = AlphaZero.player_temperature(player, game, length(trace))
    π_sample = apply_temperature(π_target, τ)
    a = actions[Util.rand_categorical(π_sample)]
    GI.play!(game, a)
    push!(trace, π_target, GI.white_reward(game), GI.current_state(game), a, rootvalue) # pᵏ, rᵏ, sᵏ, aᵏ, vᵏ
  end
end


text:
using XLA
using TensorFlow
using Test

using Flux
using Zygote

struct Relu{T}
    m::T
end
(r::Relu)(x) = relu.(r.m(x))

xrt_server_process = run(`$(joinpath(dirname(pathof(TensorFlow)),"..","deps","downloads","bin","xrt_server"))`; wait=false)
sess = Session(Graph(); target="grpc://localhost:8470")
try
    # Issue #5
    f() = XLA.HloRng(Float32, (5,5), 1)(XRTArray(0f0),XRTArray(1f0))
    @test isa(begin
        run(@tpu_compile f())
    end, XRTArray)

    # Test convolutions with respect to the weights
    model = Relu(Conv(XRTArray(fill(-1f0, 1, 1, 1, 1)), XRTArray(Float32[0f0])))
    function compute_updates(xrtic, x)
      loss, back =  let xrtic=xrtic, x=x
        Zygote._forward(
          Zygote.Context{Nothing}(nothing),
          xrtic -> sum(xrtic(x)),
          xrtic,
        )
      end
      Zygote.tailmemaybe(back(1f0))
    end
    x = reshape(Float32[1.0 -2.0; -3.0 -4.0], (2, 2, 1, 1))
    let compld = @tpu_compile compute_updates(model, XRTArray(x))
        res = run(compld, XRTRemoteStruct(sess, model), XRTArray(x))
        res = convert(typeof(res).parameters[1], res)
        @test convert(Array, res[1].m.weight)[] == -9f0
    end

    # Bringing back allocations
    f(a, b) = (a, b)
    A, B = XRTArray(sess, rand(Float32, 10, 10)), XRTArray(sess, rand(Float32, 10, 10))
    result = run((@tpu_compile f(A, B)), A, B)
    @test isa(result, XLA.XRTRemoteStruct)
    @test isa(fetch(result), Tuple{XRTArray{Float32, (10, 10), 2}, XRTArray{Float32, (10, 10), 2}})

    # Test scalar broadcasting
    bcastfoo() = Base.Broadcast.broadcasted(+, XRTArray(1), XRTArray(1))[]
    run(@tpu_compile bcastfoo())

    # Test getindex with vector (for OneHotMatrix)
    from = Float32.(reshape(1:(32*10), (32, 10)))
    idxs = [1, 2, 3, 2, 1]
    f_getidx(from, idxs) = from[:, idxs]
    let compld = @tpu_compile f_getidx(XRTArray(from), XRTArray(idxs))
        @test convert(Array, run(compld, XRTArray(from), XRTArray(idxs))) ==
            f_getidx(from, idxs)
    end

    arr = Float32[1.0 2.0; 3.0 4.0]
    f(A) = repeat(arr, 2, 2)
    let compld = @tpu_compile f(XRTArray(arr))
        @test convert(Array, run(compld, XRTArray(arr))) == f(arr)
    end
    vec = Float32[1.0 2.0]
    let compld = @tpu_compile f(XRTArray(vec))
        @test convert(Array, run(compld, XRTArray(vec))) == f(vec)
    end

finally
    GC.gc() # Try running finalizers, before the process exits
    close(sess)
    kill(xrt_server_process)
end


text:
using DotfileManager
using Test

@testset "ArgParse" begin
	@test DotfileManager.parse_commandline() == Dict(["parse" => nothing])
end