text:
module SelfPhase
export derive_constants, simulate, initialize, saveData, PropagationError

using Dierckx
using SpecialFunctions.dawson
using Base.Filesystem
using ProgressMeter
using YAML

#=
 ██████  ██████  ███    ██ ███████ ████████  █████  ███    ██ ████████ ███████
██      ██    ██ ████   ██ ██         ██    ██   ██ ████   ██    ██    ██
██      ██    ██ ██ ██  ██ ███████    ██    ███████ ██ ██  ██    ██    ███████
██      ██    ██ ██  ██ ██      ██    ██    ██   ██ ██  ██ ██    ██         ██
 ██████  ██████  ██   ████ ███████    ██    ██   ██ ██   ████    ██    ███████
=#
# Does this speed things up??
FFTW.set_num_threads(nprocs())
BLAS.set_num_threads(nprocs())

struct PropagationError <: Exception
    msg::String
end

###################
# Fixed Constants #
###################
# Setup constants
const losses  = 0.13          # Absorption Coef.        1/m
const Ui_Ar   = 15.75*1.6E-19 # Ionization Energy of Ar J
const α       = 7E-13         # ???
const Zeff_Ar = 1.0           # Effective Charge of ionized Argon

# General Physics Constants
const c  = 299792458.0      # Speed of light     m/s
const ħ = 1.0545718E-34     # Reduced Planck     Js
const me = 9.10938356E-31   # Electron Mass      kg
const ee = 1.6021766208E-19 # Elementary Charge  C
const ϵ0 = 8.854187817E-12  # Vaccum Permitivity F/m
const Kb = 1.38064852E-23   # Boltzmann Constant J/K
const T  = 300.0            # ~Room Temperature  K

# Indices of refraction/dispersion for Argon 1/nm
const C1 = 0.012055
const C2 = 0.2075
const C3 = 91.012
const C4 = 0.0415
const C5 = 87.892
const C6 = 4.3330
const C7 = 214.02
C = [C1,C2,C3,C4,C5,C6,C7]

#Indices of Refraction for Fused Silica 1/nm
const Cfs1 = 0.6961663
const Cfs2 = 0.0684043
const Cfs3 = 0.4079426
const Cfs4 = 0.1162414
const Cfs5 = 0.8974794
const Cfs6 = 9.896161
Cfs = [Cfs1,Cfs2,Cfs3,Cfs4,Cfs5,Cfs6]

# Plasma Ionization Constants
const a0 = -185.8
const a1 =  11.16
const a2 =  4.763e-3
const a3 = -9.946e-4
const a4 = -9.722e-5
const a5 =  1.182e-5
const a6 =  4.63e-7
const a7 = -4.227e-8
as = [a0,a1,a2,a3,a4,a5,a6,a7]

#=
███████ ██    ██ ███    ██  ██████ ████████ ██  ██████  ███    ██ ███████
██      ██    ██ ████   ██ ██         ██    ██ ██    ██ ████   ██ ██
█████   ██    ██ ██ ██  ██ ██         ██    ██ ██    ██ ██ ██  ██ ███████
██      ██    ██ ██  ██ ██ ██         ██    ██ ██    ██ ██  ██ ██      ██
██       ██████  ██   ████  ██████    ██    ██  ██████  ██   ████ ███████
=#

#################
# File Handling #
#################
"""
    saveData(fname, E, ΔT_pulse, z)

Saves all the data to respective files within the folder given by fname
"""
function saveData(fname, E, ΔT_pulse, z)
    open(fname * "/E", "a") do f
        #write(f, @sprintf("\# z = %f\n", z))
        for pair in zip(real(E),imag(E))
            write(f, @sprintf("%.15e,%.15e\n",pair...))
        end
        write(f, "\n")
    end
    open(fname * "/z", "a") do f
        write(f, "$z\n")
    end
end

"""
    loadParams(fname)

Loads the parameters from fname, a yaml file, into a dictionary.
"""
function loadParams(fname)
    p = YAML.load(open(fname))

    # We use λ in the code, so replace the key by removing the value and reinserting it
    if haskey(p, "lambda")
        merge!(p, Dict("λ" => pop!(p, "lambda")))
    end

    # Derive additional constants
    derive_constants(p)
    return p
end

"""
    saveParams(fname, p)

Saves the important parameters from a dictionary into a file given by fname.
"""
function saveParams(fname, p)
    open("$fname/params", "w") do f
        write(f, @sprintf("Energy:    [%e]\n", p["Energy"]))
        write(f, @sprintf("Tfwhm:     [%e]\n", p["Tfwhm"]))
        write(f, @sprintf("lambda:    [%e]\n", p["λ"]))
        write(f, @sprintf("dz:        [%f]\n", p["dz"]))
        write(f, @sprintf("zmax:      [%f]\n", p["zmax"]))
        write(f, @sprintf("Nt:        [%d]\n", p["Nt"]))
        write(f, @sprintf("tmax:      [%e]\n", p["tmax"]))
        write(f, @sprintf("Pin:       [%f]\n", p["Pin"]))
        write(f, @sprintf("Pout:      [%f]\n", p["Pout"]))
        write(f, @sprintf("fiberD:    [%e]\n", p["fiberD"]))
        if haskey(p, "Chirp")
            write(f, @sprintf("Chirp:     [%d]\n",     p["Chirp"]))
        end
        if haskey(p, "TOD")
            write(f, @sprintf("TOD:       [%d]\n",     p["TOD"]))
        end
    end
    writedlm("$fname/t_vec", p["t_vec"])
end

"""
    initialize(fname, p, resume, keep)

Initializes a simulation given a folder name that may or may not already contain
data. Folder name should be relevant to the paramters in p. If resume is true,
will attempt to continue a simulaiton. If keep is true, will not overwrite old
data already present in folder given by fname.
"""
function initialize(fname, p, resume, keep)
    zinit = 0
    # Check for existing folder
    if fname ∈ readdir()
        # Resume simulations if desired
        if resume
            zinit = float(read("$fname/z"))
            E[:] = readcsv("$fname/E")[end][:]
        else
            # Don't overwrite data, append number to folder name.
            if keep
                i = 2
                while fname*"_($i)" ∈ readdir()
                    i += 1
                end
                fname = fname*"_($i)"
            else
                rm(fname, recursive=true)
            end
            # Otherwise make new folder and initialize data
            mkdir(fname)
            E = initField(p)
            saveParams(fname, p)
        end
    else
        # Otherwise make new folder and initialize data
        mkdir(fname)
        E = initField(p)
        saveParams(fname, p)
    end
    return E, zinit
end

#######################
# Handling Parameters #
#######################
"""
    derive_constants(p)

Adds several concrete values and vectors to the parameter dictionary.
"""
function derive_constants(p)
    if haskey(p, "lambda")
        merge!(p, Dict("λ" => pop!(p, "lambda")))
    end

    f      = c / p["λ"]                 # Pulse Frequency Hz
    ω      = 2*pi*f                       # Pulse Angular Frequency
    σ_t    = p["Tfwhm"]/sqrt(2*log(2))    # 1-sigma width of pulse
    Power  = sqrt(2/pi) * p["Energy"]/σ_t # Max power delivered by pulse

    dt     = p["tmax"]/(p["Nt"]-1)        # Time step
    points = (-p["Nt"]/2:1:p["Nt"]/2-1)
    t_vec  = (-p["Nt"]/2:1:p["Nt"]/2-1)*dt  # Time grid iterator

    ff     = points./p["tmax"]            # Frequency grid
    ωω     = (2*pi)*ff                    # Angular frequency grid

    # Peak centered Wavelength grid, in nm
    λ_tot  = 1E9 * c ./ (f + ff)
    λ_tot_micron = 1E6 * c ./ (f + ff)
    # Peak centered angular frequency grid
    ωω_tot = ω+ωω

    # Nonlinear index of refraction
    n_tot_0 = 1 + C1 * (C2 * (λ_tot_micron.^2) ./ (C3 * (λ_tot_micron.^2) -1) +
                        C4 * (λ_tot_micron.^2) ./ (C5 * (λ_tot_micron.^2) -1) +
                        C6 * (λ_tot_micron.^2) ./ (C7 * (λ_tot_micron.^2) -1))

    λ_min = 400
    λ_max = 1500
    λ_min = λ_tot[indmin(abs.(λ_tot-λ_min))]
    λ_max = λ_tot[indmin(abs.(λ_tot-λ_max))]

    ρ_crit = ω^2*me*ϵ0/ee^2
    k_Ar   = ceil(Ui_Ar / (ħ*ω))

    dp = Dict(
        "f"       => f,
        "ω"       => ω,
        "σ_t"     => σ_t,
        "Power"   => Power,

        "dt"      => dt,
        "points"  => points,
        "t_vec"   => t_vec,

        "ff"      => ff,
        "ωω"      => ωω,

        "λ_tot"   => λ_tot,
        "ωω_tot"  => ωω_tot,
        "λ_max"   => λ_max,
        "λ_min"   => λ_min,

        "n_tot_0" => n_tot_0,

        "ρ_crit"  => ρ_crit,
        "k_Ar"    => k_Ar
    )
    merge!(p,dp)
end

######################
# Simulation Helpers #
######################
"""
    initField(p)

Returns a vector representing an electric field defined in terms of the
peak power and spread of the intial gaussian pulse. Can also handle
Chirp and TOD. All paramters read from dictionary p.
"""
function initField(p)
    E             = exp.(-p["t_vec"].^2/p["σ_t"]^2)
    E0            = sqrt(2*p["Power"]/(pi*(p["fiberD"]/2)^2))
    E             = E0 .* E

    # Use get so as to have a default value of zero instead
    # of an exception in case of Chirp or TOD not being defined.
    Chirp_function = get(p,"Chirp", 0) * p["ωω"] .^ 2 +
                     get(p, "TOD", 0) .* p["ωω"] .^ 3
    if !iszero(Chirp_function)
	       E_TF           = fftshift(fft(fftshift(E))).*exp.(im * Chirp_function)
           E              = ifftshift(ifft(ifftshift(E_TF)))
    end

    return E
end

"""
    calc_pressure(p_in, p_out, z, zmax)

Computes the pressure at a point z along a linear pressure gradient.
"""
function calc_pressure(p_in, p_out, z, zmax) #tested
    #Calculate pressure along fiber
    return sqrt(p_in^2 + (z/zmax) * (p_out^2-p_in^2))
end

"""
    calc_duration(E, t1)

Computes the FHWM of a gaussian peak, doesn't appear to work properly.
"""
function calc_duration(E, t1) #tested
    center_pulse = sum(t1.*(abs2.(E)))/sum(abs2.(E))
    return 2 * sqrt(2*log(2)).*((sum((t1-center_pulse).^2.*abs2.(E))
                               / sum(abs.(E).^2)).^0.5)*1E15
end

"""
    prop_lin(p, E, deriv_t_2, losses, ft, ift)

Computes the linear evoltion of a field E. deriv_t_2 is computed in sim step.
ft and ift are pre-planned FFT matrices for forward and reverse FT. losses is
the absorption coeff of the fiber.
"""
function prop_lin(p, E, deriv_t_2, losses, ft, ift) #tested
    # Shift to frequency domain and compute linear propagation
    E_TF = fftshift(ft * (fftshift(E))) .* exp.(1im*(deriv_t_2)* p["dz"])
    # Shift back to time domain, compute losses and return
    return ifftshift(ift * (ifftshift(E_TF))) .* exp.(-losses/2 * p["dz"])
end

"""
    prop_non_lin(p, E, rrr, ρ, losses, kerr_response)

Computes the nonlinear evolution of a field E. rrr, ρ, losses and kerr_response
all computed in main.
"""
function prop_non_lin(p, E, rrr, ρ, losses, kerr_response) #tested
    return E.*exp.(rrr.*ρ*p["dz"] - losses + kerr_response)
end

"""
    smooth(values, radius)

Smooths a vector by averaging all values within radius of each element.
"""
function smooth(values, radius)
    smoothed_values = zeros(values)
    for i in eachindex(values)
        # Handles edges by smoothing over a smaller radius
        temp_radius = minimum([radius, i - 1, length(values) - i])
        smoothed_values[i] = mean(values[(i - temp_radius):(i + temp_radius)])
    end
    return smoothed_values
end

"""
    NL_response(p, E, γs)

Computes the electric field brought on by steepening.
"""
function NL_response(p, E, γs)
    #=
    Functionized version of the repeated code in steppening.m
    =#
    NL = sum([γs[i] * abs.(E).^(2*i) for i in eachindex(γs)]) * p["dz"]
    Temp = NL .* E
    return (im / p["ω"]) *
           ((circshift(Temp,1) - circshift(Temp,-1))/(2 * p["dt"]))
end

"""
    steepening(p, E, γs)

Computes the evoltion of field E due to steepening.
"""
function steepening(p, E, γs) #Tested
    Etemp = E + (0.5 * NL_response(p, E, γs))
    return E + NL_response(p, Etemp, γs)
end

"""
    calc_ks(p, n_tot)

Computes the kerr coefficients for a given frequency.
"""
function calc_ks(p, n_tot) #tested
    # Interpolation of n
    n_interp = Spline1D(p["ωω_tot"], n_tot,k=4)
    n_tot = evaluate(n_interp, p["ωω_tot"])
    dn    = derivative(n_interp, p["ωω_tot"], nu=1)
    d2n   = derivative(n_interp, p["ωω_tot"], nu=2)
    d3n   = derivative(n_interp, p["ωω_tot"], nu=3)
    d4n   = derivative(n_interp, p["ωω_tot"], nu=4)

    # Frequency Disperion
    k_tot = n_tot .* p["ωω_tot"]/c
    k_tot[p["λ_tot"] .< 245] = maximum(k_tot)

    k_first  = 1/c * (dn .* p["ωω_tot"] + n_tot)
    k_second = 1/c * (d2n .* p["ωω_tot"] + 2 * dn)
    k_third  = 1/c * (d3n .* p["ωω_tot"] + 3 * d2n)
    k_fourth = smooth(1/c * (d4n .* p["ωω_tot"] + 4 * d3n), 100)

    k  = k_tot[findfirst(x->x==p["ω"],p["ωω_tot"])]
    k1 = k_first[findfirst(x->x==p["ω"],p["ωω_tot"])]
    k2 = k_second[findfirst(x->x==p["ω"],p["ωω_tot"])]
    k3 = k_third[findfirst(x->x==p["ω"],p["ωω_tot"])]
    k4 = k_fourth[findfirst(x->x==p["ω"],p["ωω_tot"])]
    return k_tot, [k,k1,k2,k3,k4]
end

"""
    calc_ns(pressure, n, n_tot, λ_tot)

Computes the index of refraction for a given pressure along the entire frequency
range.
"""
function calc_ns(pressure, n, n_tot, λ_tot) #tested
    n2_800  = 1e-23   * pressure
    n4_800  =-3.7e-42 * pressure
    n6_800  = 4e-58   * pressure
    n8_800  =-1.7e-75 * pressure
    n10_800 = 8.8e-94 * pressure

    n_800 = n_tot[indmin(abs.(λ_tot - 800))]
    n2  = n2_800 * ((n^2 - 1)/(n_800^2-1))^4
    n4  = n4_800 * ((n^2 - 1)/(n_800^2-1))^6
    n6  = n6_800 * ((n^2 - 1)/(n_800^2-1))^8
    n8  = n8_800 * ((n^2 - 1)/(n_800^2-1))^10
    n10 = n10_800 * ((n^2 - 1)/(n_800^2-1))^12

    return [n,n2,n4,n6,n8,n10]
end

"""
    plasma_potential(E, p, Ui)

Computes the chance of ionization based off of PPT theory.
"""
function plasma_potential(E,p,Ui) #Tested
    """
    Derived From PPT Theory
    http://jetp.ac.ru/cgi-bin/dn/e_023_05_0924.pdf
    """
    Zeff = 1
    Uh = 13.5984*ee                     # Hydrogen Ionization Potential
    ω_au = 4.1E16                       # Ionization Potential (1/s Natural)
    E = abs.(E) * sqrt(2/(ϵ0*c))        # Renomrmalize E-field for proper units
    γ = p["ω"] .* sqrt(2*me*Ui)./(ee*E)
    Eh = ee^5*me^2/(ħ^4 * (4*π*ϵ0)^3)
    E0 = Eh * (Ui/Uh) ^ (3/2)
    A = zeros(γ)
    β = 2*γ ./ sqrt.(1+γ.^2)
    α = 2.*asinh.(γ) - β

    g=3./(2*γ).*((1+1./(2*γ.^2)).*asinh.(γ)-1./β)
    ν0 = Ui / (ħ*p["ω"])
    ν  = ν0 * (1 + 1./(2*γ.^2))
    kmin = minimum(floor.(ν) + 1)
    # These Values are gas dependent, see paper.
    l=0
    m=0
    n_star = Zeff * sqrt(Uh/Ui)

    C_nl2 = 2^(2*n_star) / (n_star*gamma(2*n_star))
    f = (2*l+1) * factorial(l+abs(m)) / (2^abs(m)) *
        factorial(abs(m)) * factorial(l-abs(m))

    A = sum([4/sqrt(3*π) * γ.^2 ./ (1+γ.^2) .* exp.(-α .* (z-ν)) .*
             dawson.(sqrt.(complex(abs.(β.*(z - ν))))) for z in kmin:kmin+2])

    potential = ω_au * C_nl2 * f * sqrt(6/π) *
                (Ui / (2 * Uh)) * A .*
                (2 * E0./(E .* sqrt.(1+γ.^2))).^(2 * n_star - abs(m) - 3/2) .*
                exp.(-2 * E0 * g ./ (3*E))

    # present in original code, seemingly not needed
    # potential[isnan.(potential)]=0

    return potential
end

"""
    plasma(p, α, ρ_at, Potentiel_Ar, E, coeff2)

Computes the plasma density along the pulse.
"""
function plasma(p, α, ρ_at, Potentiel_Ar, E, coeff2) #tested
    ρ_Ar = zeros(size(E))
    for i in 1:(p["Nt"]-1)
        (ρ_Ar[i+1] = ρ_Ar[i] +
                    p["dt"] * (-α*ρ_Ar[i]^2 + Potentiel_Ar[i]*(ρ_at - ρ_Ar[i]) +
                               (coeff2 * abs(E[i])^2)*ρ_Ar[i])
        ) |> isfinite || throw(PropagationError("Divergence in plasma calculation."))
    end
    return ρ_Ar
end

"""
    simulate(E, p, zinit, fname, num_saves)

Performs the full simuation of a pulse propagating. Saves results to fname. Will
save num_saves times with even propagation length between saves.
"""
function simulate(E, p, zinit, fname, num_saves)

    ##################
    # Initialisation #
    ##################
    nsteps = ceil(Int, (p["zmax"]-zinit)/p["dz"])

    # How often to save data.
    if num_saves > 2
        save_every = floor(Int, nsteps/num_saves)
    else
        # Data is explicitely saved at initial and final points.
        save_every = Inf
    end

    #Plan FFT
    ft  = plan_fft(E)
    ift = plan_ifft(E)

    # Propagation variables
    z = zinit
    ρ = 0
    for i=0:nsteps-1
        if (i % save_every == 0)
            # Async data write
            saveData(fname, E,
                     calc_duration(E,p["t_vec"]*p["dt"]), z)
        end
        open(fname * "/PlasmaDensity", "a") do f
            write(f, @sprintf("%.5e\n", ρ))
        end

        E, ρ = simStep(E, p, z, ft, ift)

        # Update Distance
        z += p["dz"]
    end

    #Save final data
    saveData(fname, E, calc_duration(E,p["t_vec"]*p["dt"]), z)
end

"""
    simStep(E, p, z, ft, ift)

Performs a single split time step of the simulation.
"""
function simStep(E, p, z, ft, ift)
    # Calculate Pressure
    pressure_z = calc_pressure(p["Pin"], p["Pout"], z, p["zmax"])

    # Update of n, with cutoffs
    n_tot = sqrt.(complex(1+pressure_z * (p["n_tot_0"].^2 - 1)))
    n_max = n_tot[findfirst(λ->λ==p["λ_max"], p["λ_tot"])]
    n_tot[p["λ_tot"] .> p["λ_max"]] = n_max
    n_min = n_tot[findfirst(λ->λ==p["λ_min"], p["λ_tot"])]
    n_tot[p["λ_tot"] .< p["λ_min"]] = n_min

    k_tot, ks = calc_ks(p, n_tot)
    n  = n_tot[findfirst(x->x==p["ω"],p["ωω_tot"])]
    vg = 1/ks[2]

    # Argon Parameters
    ns = calc_ns(pressure_z, n, n_tot, p["λ_tot"])
    #β2 = pressure_z * ks[3]
    #β3 = pressure_z * ks[4]
    #β4 = pressure_z * ks[5]
    τ = 3.5E-13 / pressure_z
    ρ_at = pressure_z * 1E5 / (Kb * T)

    # Plasma Parameters
    σ_k = 2.81E-96 * p["Pout"]
    σ   = (ks[1]*ee^2) ./ (p["ω"] * me * ϵ0) .* τ./(1+(p["ω"] * τ).^2)
    #β_k = 10.^(-4 * p["k_Ar"]) .* p["k_Ar"] * ħ .* p["ω"] * ρ_at * 0.21 * σ_k
    rrr = -im * ks[1]./(2 * n[1]^2 * p["ρ_crit"]) - 0.5 * σ
    coeff2 = σ/Ui_Ar

    # Dispersion and Laplacian Operators
    dv_t_2_op = k_tot - ks[1] - ks[2] * p["ωω"]

    # Kerr factors
    γs = im * ns[2:end] * ks[1]/n[1]

    # Propagation
    E = prop_lin(p, E, -dv_t_2_op, losses, ft, ift)   #Linear
    E = steepening(p, E, γs)                          #Steepening

    # Plasma
    U_ion = plasma_potential(E, p, Ui_Ar)
    ρ = plasma(p, α, ρ_at, U_ion, E, coeff2)
    plasma_loss = U_ion ./ (2 * abs.(E).^2) * Ui_Ar .* (ρ_at - ρ) * p["dz"]
    plasma_loss[isnan.(plasma_loss)] = 0

    # Kerr and Plasma Propagation (NonLinear)
    kerr_response = -(γs[1] * abs.(E).^2 +
                      γs[2] * abs.(E).^4 +
                      γs[3] * abs.(E).^6 +
                      γs[4] * abs.(E).^8 +
                      γs[5] * abs.(E).^10
                     ) * p["dz"]

    E = prop_non_lin(p, E, rrr, ρ, plasma_loss, kerr_response)

    return E, maximum(ρ) * 1E-6
end

#module
end

max_stars_repo_name: ry-dgel/self-phase

max_issues_repo_name: ry-dgel/self-phase

max_forks_repo_name: ry-dgel/self-phase


text:
"""
    stacktraces(in; <keyword arguments>)

Stack an array of traces
# Arguments
- `in`: 2D seismic data
- `normalize=true`: Normalize result by number of traces

```
"""
function stacktraces(d;normalize=true)
	if (normalize == true)
		val = sum(d[:,:],dims=2)/size(d[:,:],2)
	else
		val = sum(d[:,:],dims=2)
	end
	return val
end

ext: jl


text:
"""
    SGHMC{AD,space}

Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) sampler.e

# Fields
$(TYPEDFIELDS)

# Reference

Tianqi Chen, Emily Fox, & Carlos Guestrin (2014). Stochastic Gradient Hamiltonian Monte
Carlo. In: Proceedings of the 31st International Conference on Machine Learning
(pp. 1683–1691).
"""
struct SGHMC{AD,space,T<:Real} <: StaticHamiltonian{AD}
    learning_rate::T
    momentum_decay::T
end

"""
    SGHMC{AD}(space::Symbol...; learning_rate, momentum_decay)

Create a Stochastic Gradient Hamiltonian Monte Carlo (SGHMC) sampler.

If the automatic differentiation backend `AD` is not provided, the currently activated
AD backend in Turing is used.

# Reference

Tianqi Chen, Emily Fox, & Carlos Guestrin (2014). Stochastic Gradient Hamiltonian Monte
Carlo. In: Proceedings of the 31st International Conference on Machine Learning
(pp. 1683–1691).
"""
SGHMC(args...; kwargs...) = SGHMC{ADBackend()}(args...; kwargs...)
function SGHMC{AD}(
    space::Symbol...;
    learning_rate::Real,
    momentum_decay::Real,
) where {AD}
    _learning_rate, _momentum_decay = promote(learning_rate, momentum_decay)
    return SGHMC{AD,space,typeof(_learning_rate)}(_learning_rate, _momentum_decay)
end

struct SGHMCState{V<:AbstractVarInfo, T<:AbstractVector{<:Real}}
    vi::V
    velocity::T
end

function DynamicPPL.initialstep(
    rng::Random.AbstractRNG,
    model::Model,
    spl::Sampler{<:SGHMC},
    vi::AbstractVarInfo;
    kwargs...,
)
    # Transform the samples to unconstrained space and compute the joint log probability.
    if !DynamicPPL.islinked(vi, spl)
        DynamicPPL.link!(vi, spl)
        model(rng, vi, spl)
    end

    # Compute initial sample and state.
    sample = Transition(vi)
    state = SGHMCState(vi, zero(vi[spl]))

    return sample, state
end

function AbstractMCMC.step(
    rng::Random.AbstractRNG,
    model::Model,
    spl::Sampler{<:SGHMC},
    state::SGHMCState;
    kwargs...
)
    # Compute gradient of log density.
    vi = state.vi
    θ = vi[spl]
    _, grad = gradient_logp(θ, vi, model, spl)

    # Update latent variables and velocity according to
    # equation (15) of Chen et al. (2014)
    v = state.velocity
    θ .+= v
    η = spl.alg.learning_rate
    α = spl.alg.momentum_decay
    newv = (1 - α) .* v .+ η .* grad .+ sqrt(2 * η * α) .* randn(rng, eltype(v), length(v))

    # Save new variables and recompute log density.
    vi[spl] = θ
    model(rng, vi, spl)

    # Compute next sample and state.
    sample = Transition(vi)
    newstate = SGHMCState(vi, newv)

    return sample, newstate
end

"""
    SGLD

Stochastic gradient Langevin dynamics (SGLD) sampler.

# Fields
$(TYPEDFIELDS)

# Reference

Max Welling & Yee Whye Teh (2011). Bayesian Learning via Stochastic Gradient Langevin
Dynamics. In: Proceedings of the 28th International Conference on Machine Learning
(pp. 681–688).
"""
struct SGLD{AD,space,S} <: StaticHamiltonian{AD}
    "Step size function."
    stepsize::S
end

struct PolynomialStepsize{T<:Real}
    "Constant scale factor of the step size."
    a::T
    "Constant offset of the step size."
    b::T
    "Decay rate of step size in (0.5, 1]."
    γ::T

    function PolynomialStepsize{T}(a::T, b::T, γ::T) where T
        0.5 < γ ≤ 1 || error("the decay rate `γ` has to be in (0.5, 1]")
        return new{T}(a, b, γ)
    end
end

"""
    PolynomialStepsize(a[, b=0, γ=0.55])

Create a polynomially decaying stepsize function.

At iteration `t`, the step size is
```math
a (b + t)^{-γ}.
```
"""
function PolynomialStepsize(a::T, b::T, γ::T) where {T<:Real}
    return PolynomialStepsize{T}(a, b, γ)
end
function PolynomialStepsize(a::Real, b::Real = 0, γ::Real = 0.55)
    return PolynomialStepsize(promote(a, b, γ)...)
end

(f::PolynomialStepsize)(t::Int) = f.a / (t + f.b)^f.γ

"""
    SGLD{AD}(space::Symbol...; stepsize = PolynomialStepsize(0.01))

Stochastic gradient Langevin dynamics (SGLD) sampler.

By default, a polynomially decaying stepsize is used.

# Reference

Max Welling & Yee Whye Teh (2011). Bayesian Learning via Stochastic Gradient Langevin
Dynamics. In: Proceedings of the 28th International Conference on Machine Learning
(pp. 681–688).

See also: [`PolynomialStepsize`](@ref)
"""
SGLD(args...; kwargs...) = SGLD{ADBackend()}(args...; kwargs...)
function SGLD{AD}(
    space::Symbol...;
    stepsize = PolynomialStepsize(0.01),
) where {AD}
    return SGLD{AD,space,typeof(stepsize)}(stepsize)
end

struct SGLDTransition{T,F<:Real}
    "The parameters for any given sample."
    θ::T
    "The joint log probability of the sample."
    lp::F
    "The stepsize that was used to obtain the sample."
    stepsize::F
end

function SGLDTransition(vi::AbstractVarInfo, stepsize)
    theta = tonamedtuple(vi)
    lp = getlogp(vi)
    return SGLDTransition(theta, lp, stepsize)
end

metadata(t::SGLDTransition) = (lp = t.lp, SGLD_stepsize = t.stepsize)

DynamicPPL.getlogp(t::SGLDTransition) = t.lp

struct SGLDState{V<:AbstractVarInfo}
    vi::V
    step::Int
end

function DynamicPPL.initialstep(
    rng::Random.AbstractRNG,
    model::Model,
    spl::Sampler{<:SGLD},
    vi::AbstractVarInfo;
    kwargs...
)
    # Transform the samples to unconstrained space and compute the joint log probability.
    if !DynamicPPL.islinked(vi, spl)
        DynamicPPL.link!(vi, spl)
        model(rng, vi, spl)
    end

    # Create first sample and state.
    sample = SGLDTransition(vi, zero(spl.alg.stepsize(0)))
    state = SGLDState(vi, 1)

    return sample, state
end

function AbstractMCMC.step(
    rng::Random.AbstractRNG,
    model::Model,
    spl::Sampler{<:SGLD},
    state::SGLDState;
    kwargs...
)
    # Perform gradient step.
    vi = state.vi
    θ = vi[spl]
    _, grad = gradient_logp(θ, vi, model, spl)
    step = state.step
    stepsize = spl.alg.stepsize(step)
    θ .+= (stepsize / 2) .* grad .+ sqrt(stepsize) .* randn(rng, eltype(θ), length(θ))

    # Save new variables and recompute log density.
    vi[spl] = θ
    model(rng, vi, spl)

    # Compute next sample and state.
    sample = SGLDTransition(vi, stepsize)
    newstate = SGLDState(vi, state.step + 1)

    return sample, newstate
end


text:
using Test
using StaticArrays

abstract type AbstractModel end
abstract type OneLayerModel <: AbstractModel end

struct EmptyModel <: OneLayerModel end
state(m::EmptyModel, T) = @vars()

struct ScalarModel <: OneLayerModel end
state(m::ScalarModel, T) = @vars(x::T)

struct VectorModel{N} <: OneLayerModel end
state(m::VectorModel{N}, T) where {N} = @vars(x::SVector{N, T})

struct MatrixModel{N, M} <: OneLayerModel end
state(m::MatrixModel{N, M}, T) where {N, M} =
    @vars(x::SHermitianCompact{N, T, M})

abstract type TwoLayerModel <: AbstractModel end

struct CompositModel{EM, SM, VM, MM} <: TwoLayerModel
    empty_model::EM
    scalar_model::SM
    vector_model::VM
    matrix_model::MM
end
function CompositModel(
    Nv,
    N,
    M;
    empty_model = EmptyModel(),
    scalar_model = ScalarModel(),
    vector_model = VectorModel{Nv}(),
    matrix_model = MatrixModel{N, M}(),
)
    args = (empty_model, scalar_model, vector_model, matrix_model)
    return CompositModel{typeof.(args)...}(args...)
end

function state(m::CompositModel, T)
    @vars begin
        empty_model::state(m.empty_model, T)
        scalar_model::state(m.scalar_model, T)
        vector_model::state(m.vector_model, T)
        matrix_model::state(m.matrix_model, T)
    end
end

Base.@kwdef struct NTupleModel{S, V, Nv} <: OneLayerModel
    scalar_model::S = ScalarModel()
    vector_model::V = VectorModel{Nv}()
end

function NTupleModel(
    Nv;
    scalar_model::S = ScalarModel(),
    vector_model::V = VectorModel{Nv}(),
) where {S, V}
    return NTupleModel{S, V, Nv}(scalar_model, vector_model)
end

function state(m::NTupleModel, T)
    @vars begin
        scalar_model::state(m.scalar_model, T)
        vector_model::state(m.vector_model, T)
    end
end

state(m::NTuple{N, NTupleModel}, FT) where {N} =
    Tuple{ntuple(i -> state(m[i], FT), N)...}

struct NTupleContainingModel{N, NTM, VM, SM} <: TwoLayerModel
    ntuple_model::NTM
    vector_model::VM
    scalar_model::SM
end
function NTupleContainingModel(
    N,
    Nv;
    ntuple_model = ntuple(i -> NTupleModel(Nv), N),
    vector_model = VectorModel{Nv}(),
    scalar_model = ScalarModel(),
)
    args = (ntuple_model, vector_model, scalar_model)
    return NTupleContainingModel{N, typeof.(args)...}(args...)
end

function state(m::NTupleContainingModel, T)
    @vars begin
        ntuple_model::state(m.ntuple_model, T)
        vector_model::state(m.vector_model, T)
        scalar_model::state(m.scalar_model, T)
    end
end


text:
# test1.jl
# based on ex1.jl
using Distributions

function test1()
println("# Starting Test 1")
num_data_default = 200
gen_data_normal(theta::Array, n::Integer = num_data_default) = rand(Normal(theta[1],theta[2]),num_data_default)

normalize_theta2_pos(theta::Array) =  theta[2] = abs(theta[2])
is_valid_theta2_pos(theta::Array) =  theta[2]>0.0 ? true : false

theta_true = [0.0, 1.0]
param_prior = Distributions.MvNormal(theta_true,ones(length(theta_true)))

abc_plan = abc_pmc_plan_type(gen_data_normal,ABC.calc_summary_stats_mean_var,ABC.calc_dist_max, param_prior;
                             target_epsilon=0.01,is_valid=is_valid_theta2_pos,num_max_attempt=10000);

num_param = 2
# Generate "observed" data
data_obs = abc_plan.gen_data(theta_true)
ss_obs = abc_plan.calc_summary_stats(data_obs)
println("# Setting 'true' parameters of theta= ",theta_true)
println("# Generated 'observed' data with summary statistics = ",vec(ss_obs))
println("# Starting ABC-PMC run")
pop_out = run_abc(abc_plan,ss_obs;verbose=true);

println("# Results: ",mean(pop_out.theta[1,:]), " ",mean(pop_out.theta[2,:]))
if abs(mean(pop_out.theta[1,:]))>0.05
  return false
end
if abs(mean(pop_out.theta[2,:])-1)>0.05
  return false
end
return true
end


text:
# Runs MAP Inference algorithms
using SPNetworks
import SPNetworks: leaves, isleaf, issum, isprod, IndicatorFunction, project, project2, binarize!
import SPNetworks.MAP: maxproduct!, localsearch!, beliefpropagation!, treebeliefpropagation!

if length(ARGS) < 2
    println("Usage: julia --color=yes mapinference.jl spn_filename query_filename [maxinstances] [solution_filename]")
    exit()
end
spn_filename = ARGS[1]
q_filename = ARGS[2]
maxinstances = 100
if length(ARGS) > 2
    maxinstances = parse(Int, ARGS[3])
end
o_filename = q_filename * ".solution"
if length(ARGS) > 3
    o_filename = ARGS[4]
end

# Which algorithms to run?
# - mp: max-product
# - ls: Local-Search
# - bp: belief-propagation
# each algorithm uses the incumbent solution of previously ran algorithms (so order matters)
# algorithms = [:mp, :ls, :bp]
# algorithms = [:mp, :bp]
# algorithms = [:mp, :bp, :ls]
# algorithms = [:mp, :ls, :bp, :tbp]
algorithms = [:mp, :ls, :prunedbp, :tbp]
# algorithms = [:mp, :ls]
# algorithms = [:mp]

# collect results (value and runtime)
results = Dict{Symbol,Array{Float64}}()
for algo in algorithms
    results[algo] = Float64[]
end

println("SPN: ", spn_filename)
println("Query: ", q_filename)
println("MaxInstances: ", maxinstances)
println("Algorithms: ", algorithms)
println()

# Load SPN from spn file (assume 0-based indexing)
loadtime = @elapsed spn = SumProductNetwork(spn_filename; offset = 1)
println("Loaded in $(loadtime)s ", summary(spn))
nvars = length(scope(spn))
# assignment vector
x = Array{Float64}(undef, nvars)

# maxch = 0
# for node in spn
#     if !isleaf(node)
#         global maxch = max(maxch, length(node.children))
#     end
# end
# if maxch > 4
#     @warn "maximum node indegree too large: $maxch. It is highly recommend to split nodes before running this."
# end

# Load query, evidence and marginalized variables
totaltime = @elapsed open(q_filename) do io
    inst = 1
    printstyled("╮\n"; color = :red)
    open(o_filename, "w") do outio
    while !eof(io)
        # fill!(x, NaN) # reset configuration
        printstyled("├(", inst, ")\n"; color = :red)
        # Read query variables
        fields = split(readline(io))
        header = fields[1]
        @assert header == "q"
        query = Set(map(f -> (parse(Int, f)+1), fields[2:end]))
        # Read marginalized variables
        fields = split(readline(io))
        header = fields[1]
        @assert header == "m"
        marg = Set(map(f -> (parse(Int, f)+1), fields[2:end]))
        for i in marg
            x[i] = NaN
        end
        # Read evidence
        fields = split(readline(io))
        evidence = Dict{Int,Int}()
        header = fields[1]
        @assert header == "e"
        for i=2:2:length(fields)
            var = parse(Int,fields[i]) + 1
            value = parse(Int,fields[i+1]) + 1
            x[var] = value
            evidence[var] = value
        end
        # @assert length(evidence) == nevid
        # println(
        #     join(
        #         setdiff(
        #             collect(1:nvars),
        #             union(query,keys(evidence))
        #             ) .- 1,
        #         " "
        #         )
        # )
        @assert (length(query) + length(marg) + length(evidence)) == nvars
        best = -Inf
        solution = similar(x)
        for algo in algorithms
            printstyled("│\t"; color = :red)
            if algo == :mp
                # Run max-product
                runtime = @elapsed maxproduct!(x, spn, query)
                printstyled("MaxProduct: "; color = :green)
            elseif algo == :ls
                # First prune network (discard constant and marginalized subnetworks)
                print("Pruning spn...")
                runtime = @elapsed spn2 = project(spn, query, x)
                println("done: $(runtime)s")
                # Run local search
                runtime += @elapsed localsearch!(x, spn2, query, 100)
                printstyled("LocalSearch: "; color = :green)
            elseif algo == :prunedbp
                # Run hybrid belief propagation with pruned network
                print("Pruning spn...")
                runtime = @elapsed spn2 = project2(spn, query, x)
                # project2 makes network nonbinary, and creates large potentials!
                # Fix no-binary nodes introduced
                runtime += @elapsed binarize!(spn2)
                println("done: $(runtime)s")
                runtime += @elapsed beliefpropagation!(x, spn2, query;
                    maxiterations = 10,
                    lowerbound = true,
                    verbose = true,
                    warmstart = true,
                    rndminit = false)
                # printstyled("BeliefPropagation: "; color = :green)
                runtime += @elapsed localsearch!(x, spn2, query, 100)
                printstyled("Pruning+BeliefPropagation+LocalSearch: "; color = :green)
            elseif algo == :bp
                # Run hybrid belief propagation
                runtime = @elapsed beliefpropagation!(x, spn, query;
                    maxiterations = 10,
                    lowerbound = true,
                    verbose = true,
                    warmstart = true,
                    rndminit = false)
                # printstyled("BeliefPropagation: "; color = :green)
                # First prune network (discard constant and marginalized subnetworks)
                print("Pruning spn...")
                runtime = @elapsed spn2 = project(spn, query, x)
                println("done: $(runtime)s")
                runtime += @elapsed localsearch!(x, spn2, query, 100)
                printstyled("BeliefPropagation+LocalSearch: "; color = :green)
                # for i = 1:30
                #     bptime = @elapsed beliefpropagation!(x, spn, query;
                #              maxiterations = 3,
                #              lowerbound = true,
                #              warmstart = false,
                #              verbose = false,
                #              rndminit = true)
                #     value = spn(x)
                #     print("$value")
                #     printstyled("  [$(bptime)s]\n"; color = :light_black)
                # end
                # runtime = @elapsed beliefpropagation!(x, spn, query;
                # maxiterations = 5,
                # lowerbound = true,
                # warmstart = true,
                # rndminit = false)
                # printstyled("BeliefPropagation: "; color = :green)
            elseif algo == :tbp
                # First prune network (discard constant and marginalized subnetworks)
                print("Pruning spn...")
                runtime = @elapsed spn2 = project(spn, query, x)
                println("done: $(runtime)s")
                # Run hybrid belief propagation (numerically unstable but more efficient)
                runtime += @elapsed treebeliefpropagation!(x, spn, query;
                    maxiterations = 10,
                    lowerbound = true,
                    verbose = true,
                    warmstart = false,
                    rndminit = false)
                # printstyled("T-BeliefPropagation: "; color = :green)
                runtime += @elapsed localsearch!(x, spn, query, 100)
                printstyled("T-BeliefPropagation+LocalSearch: "; color = :green)
            end
            value = spn(x)
            print("$value")
            printstyled("  [$(runtime)s]\n"; color = :light_black)
            if value >  best
                best = value
                solution .= x
            end
            # push!(results[algo], (value,runtime))
            push!(results[algo], value)
        end
        # save best solution to file
        println(outio, "# Instance: $inst  MAP Value: $best")
        for i = 1:length(solution)
            print(outio, solution[i], " ")
        end
        println(outio)
        inst += 1
        # if maximum no. of instances is reached, stop
        if inst > maxinstances
            break
        end
    end # end of while
    end # end of open file (for writing)
    printstyled("╯\n"; color = :red)
end
# display sumary of results
insts = length(results[algorithms[1]])
println("Total time: $(totaltime)s\nTotal instances: $insts")
cpad(s, n::Integer, p=" ") = rpad(lpad(s,div(n+length(s),2),p),n,p) # for printing centered
columns = [:it]
widths = [5]
for algo in algorithms
    push!(columns,algo)
    push!(widths,25)
end
bordercolor = :light_cyan
headercolor = :cyan
fieldcolor = :normal
printstyled('╭', join(map(w -> repeat('─',w), widths), '┬'), '╮', '\n' ; color = bordercolor)
for (col,w) in zip(columns,widths)
    printstyled("│"; color = bordercolor)
    printstyled(cpad(string(col), w); color = headercolor)
end
printstyled("│\n"; color = bordercolor)
printstyled('├', join(map(w -> repeat('─',w), widths), '┼'), '┤', '\n' ; color = bordercolor)
for it = 1:insts
    values = Any[ it ; [ results[algo][it] for algo in columns[2:end] ] ]
    for (col,width) in zip(values, widths)
        printstyled("│"; color = bordercolor)
        printstyled(lpad(col, width-1), ' '; color = fieldcolor)
    end
    printstyled("│\n"; color = bordercolor)
end
printstyled('├', join(map(w -> repeat('─',w), widths), '┼'), '┤', '\n' ; color = bordercolor)
values = Any[ "avg" ; [ sum(results[algo])/insts for algo in columns[2:end] ] ]
for (col,width) in zip(values, widths)
    printstyled("│"; color = bordercolor)
    printstyled(lpad(col, width-1), ' '; color = headercolor)
end
printstyled("│\n"; color = bordercolor)
printstyled('╰', join(map(w -> repeat('─',w), widths), '┴'), '╯', '\n' ; color = bordercolor)


text:
# StudentT distribution

export StudentT, HalfStudentT

@parameterized StudentT(ν)

@kwstruct StudentT(ν)
@kwstruct StudentT(ν,μ)
@kwstruct StudentT(ν,σ)
@kwstruct StudentT(ν,μ,σ)

StudentT(nt::NamedTuple{(:ν,:μ,:σ)}) = Affine(NamedTuple{(:μ,:σ)}(nt), StudentT(ν=nt.ν))
StudentT(nt::NamedTuple{(:ν,:μ,:ω)}) = Affine(NamedTuple{(:μ,:ω)}(nt), StudentT(ν=nt.ν))
StudentT(nt::NamedTuple{(:ν,:σ,)}) = Affine(NamedTuple{(:σ,)}(nt), StudentT(ν=nt.ν))
StudentT(nt::NamedTuple{(:ν,:ω,)}) = Affine(NamedTuple{(:ω,)}(nt), StudentT(ν=nt.ν))
StudentT(nt::NamedTuple{(:ν,:μ,)}) = Affine(NamedTuple{(:μ,)}(nt), StudentT(ν=nt.ν))

@affinepars StudentT

StudentT(ν, μ, σ) = StudentT((ν=ν, μ=μ, σ=σ))

@kwalias StudentT [
    df       => ν
    nu       => ν
    location => μ
    mean     => μ # Doesn't really always exist, but ok fine
    mu       => μ
    scale    => σ
    sigma    => σ
]

function logdensity(d::StudentT{(:ν,)}, x)
    ν = d.ν
    return  xlog1py((ν + 1) / (-2), x^2 / ν)
end

function basemeasure(d::StudentT{(:ν,)})
    inbounds(x) = true
    constℓ = 0.0
    varℓ() = loggamma((d.ν+1)/2) - loggamma(d.ν/2) - log(π * d.ν) / 2
    base = Lebesgue(ℝ)
    FactoredBase(inbounds, constℓ, varℓ, base)
end

TV.as(::StudentT) = asℝ

Base.rand(rng::AbstractRNG, T::Type, μ::StudentT{(:ν,)}) = rand(rng, Dists.TDist(μ.ν))

distproxy(d::StudentT{(:ν,)}) = Dists.TDist(d.ν)

@half StudentT

HalfStudentT(ν, σ) = HalfStudentT((ν=ν, σ=σ))

asparams(::Type{<:StudentT}, ::Val{:ν}) = asℝ₊


text:
using VertexEliminationOrder
using LightGraphs
using Test
using Random

ENV["JULIA_DEBUG"]=VertexEliminationOrder

G =  smallgraph("house")
@info "test treewidth of house"
@test treewidth_by_elimination!(G, collect(1:5)) == 2
@test nv(G) == 0

# tree
G = SimpleGraph{Int64}(5)
add_edge!(G, 1, 2)
add_edge!(G, 1, 3)
add_edge!(G, 3, 4)
add_edge!(G, 3, 5)
@info "test treewidth of tree"
@test treewidth_by_elimination!(G, [5, 4, 3, 2, 1]) == 1
@test nv(G) == 0

# complete graph
n = 60
G = SimpleGraph{Int64}(n)
for i in 1:(n - 1)
    for j in (i + 1):n
        add_edge!(G, i, j)
    end
end
@info "test treewidth of complete graph"
@test treewidth_by_elimination!(G, shuffle(collect(1:n))) == n - 1
@test nv(G) == 0

# unit test order_tree!
G = SimpleGraph{Int64}(10)
add_edge!(G, 1, 2)
add_edge!(G, 1, 3)
add_edge!(G, 3, 4)
add_edge!(G, 3, 5)
add_edge!(G, 4, 6)
add_edge!(G, 4, 7)
add_edge!(G, 7, 8)
add_edge!(G, 8, 9)
add_edge!(G, 8, 10)
nodes = [6, 7, 8, 9, 10, 11, 12, 23, 34, 45]
@info "testing tree_order!"
# we expect 7, 10, 11, 12, 34, 45 first (leafs) then 6, 9 then 8 (order in each block don't matter)
res = tree_order!(G, nodes)
@test  res == [7, 10, 11, 34, 45, 6, 23, 8, 12, 9]
# after this function G has no eddges
@test ne(G) == 0


text:
# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export LogWeight

"""
`LogWeight(x)`

| Set                                 |     ``\\oplus``      |  ``\\otimes``  | ``\\bar{0}`` | ``\\bar{1}`` |
|:-----------------------------------:|:--------------------:|:--------------:|:------------:|:------------:|
|``\\mathbb{R}\\cup\\{\\pm\\infty\\}``|``\\log(e^{x}+e^{y})``|     ``+``      |``-\\infty``  |   ``0``      |
"""
struct LogWeight{T <: AbstractFloat} <: Semiring
  x::T
end

LogWeight(x::Number) = LogWeight(float(x))

zero(::Type{LogWeight{T}}) where T = LogWeight{T}(T(-Inf))
one(::Type{LogWeight{T}}) where T = LogWeight{T}(zero(T))

*(a::LogWeight{T}, b::LogWeight{T}) where {T <: AbstractFloat} = LogWeight{T}(a.x + b.x)
+(a::LogWeight{T}, b::LogWeight{T}) where {T <: AbstractFloat} = LogWeight{T}(logadd(a.x,b.x))
/(a::LogWeight{T}, b::LogWeight{T}) where {T <: AbstractFloat} = LogWeight{T}(a.x - b.x)
reverse(a::LogWeight) = a
function logadd(y::T, x::T) where {T <: AbstractFloat}
  if isinf(x) return y end
  if isinf(y) return x end
  if x < y
    diff = x-y
    x = y
  else
    diff = y-x
  end
  return x + log1p(exp(diff))
end

# parsing
parse(::Type{S},str) where {T, S <: LogWeight{T}} = S(parse(T,str))

#properties
iscommulative(::Type{W}) where {W <: LogWeight} = true
isleft(::Type{W}) where {W <: LogWeight} = true
isright(::Type{W}) where {W <: LogWeight}= true
isweaklydivisible(::Type{W}) where {W <: LogWeight}= true
iscomplete(::Type{W}) where {W <: LogWeight}= true


text:
# # [Multirate Infinitesimal Step (MIS) Timestepping](@id MIS-Timestepping)

# In this tutorial, we shall explore the use of Multirate Infinitesimal Step
# (MIS) methods for the solution of nonautonomous (or non time-invariant) equations.
# For our model problem, we shall reuse the rising thermal bubble
# tutorial. See its [tutorial page](@ref Rising-Thermal-Bubble-Configuration)
# for details on the model and parameters. For the purposes of this tutorial,
# we will only run the experiment for a total of 500 simulation seconds.

using ClimateMachine
const clima_dir = dirname(dirname(pathof(ClimateMachine)));
include(joinpath(
    clima_dir,
    "tutorials",
    "Numerics",
    "TimeStepping",
    "tutorial_risingbubble_config.jl",
))

FT = Float64;

# Referencing the formulation introduced in the previous
# [Multirate RK methods tutorial](@ref Multirate-RK-Timestepping), we can
# describe Multirate Infinitesimal Step (MIS) methods by

# ```math
# \begin{align}
# v_i (0)
#   &= q^n + \sum_{j=1}^{i-1} \alpha_{ij} (Q^{(j)} - q^n) \\
# \frac{dv_i}{d\tau}
#   &= \sum_{j=1}^{i-1} \frac{\gamma_{ij}}{d_i \Delta t} (Q^{(j)} - q^n)
#     + \sum_{j=1}^i \frac{\beta_{ij}}{d_i} \mathcal{T}_S (Q^{(j)}, t + \Delta t c_i)
#     + \mathcal{T}_F(v_i, t^n +  \Delta t \tilde c_i + \frac{c_i - \tilde c_i}{d_i} \tau),
# \quad \tau \in [0, \Delta t d_i] \\
# Q^{(i)} &= v_i(\Delta t d_i),
# \end{align}
# ```
#
# where we have used the the stage values ``Q^{(i)} = v_i(\tau_i)`` as the
# solution to the _inner_ ODE problem, ``{\mathcal{T}_{s}}``
# for the slow component, and ``{\mathcal{T}_{f}}` for the fast
# one, as in the [Multirate RK methods tutorial](@ref Multirate-RK-Timestepping).

# Referencing the canonical form introduced in [Time integration](@ref
# Time-integration), both ``{\mathcal{T}_{f}}`` and ``{\mathcal{T}_{s}}``
# could be discretized either explicitly or implicitly, hence, they could
# belong to either ``\mathcal{F}(t, \boldsymbol{q})`` or ``\mathcal{G}(t, \boldsymbol{q})``
# term.
#
# The method is defined in terms of the lower-triangular matrices ``\alpha``,
# ``\beta`` and ``\gamma``, with ``d_i = \sum_j \beta_{ij}``,
# ``c_i = (I - \alpha - \gamma)^{-1} d`` and ``\tilde c = \alpha c``.
# More details can be found in [WenschKnothGalant2009](@cite) and
# [KnothWensch2014](@cite).

ode_solver = ClimateMachine.MISSolverType(;
    mis_method = MIS2,
    fast_method = LSRK144NiegemannDiehlBusch,
    nsubsteps = (40,),
)

timeend = FT(500)
CFL = FT(20)
run_simulation(ode_solver, CFL, timeend);

# The reader can compare the Courant number (denoted by `CFL` in the code snippet)
# used in this example, with the adopted in the
# [single-rate explicit timestepping tutorial page](@ref Single-rate-Explicit-Timestepping)
# in which we use the same scheme as the fast method employed in this case,
# and notice that with this MIS method we are able to take a much larger
# Courant number.

# ## References
# - [WenschKnothGalant2009](@cite)
# - [KnothWensch2014](@cite)


text:
using Choices
using Test

# from https://en.wikipedia.org/wiki/Sudoku
puzzle = [
  5 3 0 0 7 0 0 0 0;
  6 0 0 1 9 5 0 0 0;
  0 9 8 0 0 0 0 6 0;
  8 0 0 0 6 0 0 0 3;
  4 0 0 8 0 3 0 0 1;
  7 0 0 0 2 0 0 0 6;
  0 6 0 0 0 0 2 8 0;
  0 0 0 4 1 9 0 0 5;
  0 0 0 0 8 0 0 7 9;
]
solution = [
  5 3 4 6 7 8 9 1 2;
  6 7 2 1 9 5 3 4 8;
  1 9 8 3 4 2 5 6 7;
  8 5 9 7 6 1 4 2 3;
  4 2 6 8 5 3 7 9 1;
  7 1 3 9 2 4 8 5 6;
  9 6 1 5 3 7 2 8 4;
  2 8 7 4 1 9 6 3 5;
  3 4 5 2 8 6 1 7 9;
]

function block_is_consistent(block)
  have = fill(false, 9)
  for n ∈ block
    n == 0 && continue
    have[n] && return false
    have[n] = true
  end
  return true
end

"""
    cell_block(i, j)

Returns indices (a tuple of two ranges) for the 3×3 sudoku block containing the
cell at (i, j).
"""
cell_block(i, j) = map(x -> x .+ (1:3), trunc.(Int, ((i, j) .- 1) ./ 3) .* 3)

cell_is_consistent(p, i, j) =
  block_is_consistent(p[i,1:9]) &&
  block_is_consistent(p[1:9,j]) &&
  block_is_consistent(p[cell_block(i,j)...])

iscomplete(p) = !any(n == 0 for n ∈ p)

function solve(p)
  # start with the unsolved puzzle
  comp = choose([p])
  # for each "blank" (zero) cell
  for k ∈ keys(p)
    if (p[k] == 0)
      # add a piece of the computation that makes a guess for that cell and
      # checks it
      comp =
	comp |> before ->
	  choose(1:9) |> function(guess)
	    after = copy(before) # super inefficient, but whatever, it's a test
	    after[k] = guess
	    cell_is_consistent(after, Tuple(k)...) || fail()
	    return after
	  end
    end
  end
  s = first(comp)
  cut()
  return s
end

@test solve(puzzle) == solution


text:
##############################################################################
# an array of strings with header info

import Statistics

mutable struct DFrame
    A
    names
    cols
end

# empty frame
function DFrame(numrows::Integer; estnumcols=0)
    A =  Array{Float64}(undef,numrows,estnumcols)
    names = Any[]
    return DFrame(A,names,0)
end

# frame with no names
function DFrame(A)
    d = size(A,2)
    names = Array{Any}(d)
    fill!(names, "")
    return DFrame(A,names,d)
end

function DFrame(A, names)
    return DFrame(A, names, size(A,2))
end

findcolumn(c::String, DF::DFrame) = findvalue(c, DF.names)
findcolumn(c::Number, DF::DFrame) = c

function addcolumn(DF::DFrame, u::Array{Float64,2})
    newcols = size(u,2)
    if DF.cols + newcols < size(DF.A,2)
        DF.A[:,DF.cols+1:DF.cols+newcols] = u
        DF.cols += newcols
        return
    end
    DF.A = [DF.A u]
    DF.cols += newcols
end

addcolumn(DF::DFrame, u::Array{Float64,1}) = addcolumn(DF, reshape(u, length(u),1))

import Base.size
function size(DF::DFrame, x::Number)
    if x == 1
        return size(DF.A,1)
    end
    return DF.cols
end

function size(DF::DFrame)
    return size(DF.A,1), DF.cols
end
function value(DF::DFrame)
    return DF.A[:,1:DF.cols]
end

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

abstract type DataSource end

##############################################################################
# interface

mutable struct FrameSource<:DataSource
    Uf
    Vf
    Xmaps
    Ymaps
    count
    Xnames  # saved for unembedding
    Ynames  # saved for unembedding
end

function FrameSource(Uf::DFrame, Vf::DFrame)
    return FrameSource(Uf, Vf, Any[], Any[], [0], nothing, nothing)
end

function makeFrameSource(U, V, Unames, Vnames; kwargs...)
    if Unames != nothing
        Uf = DFrame(matrix(U), Unames)
    else
        Uf = DFrame(matrix(U))
    end
    if Vnames != nothing
        Vf = DFrame(matrix(V), Vnames)
    else
        Vf = DFrame(matrix(V))
    end
    return FrameSource(Uf, Vf)
end

function getXY(F::FrameSource; Uestnumcols=0, Vestnumcols=0, verbose=false)
    Xf = applyfmaplist(F.Xmaps, F.Uf; estnumcols = Uestnumcols, verbose=verbose)
    Yf = applyfmaplist(F.Ymaps, F.Vf; estnumcols = Vestnumcols, verbose=verbose)
    X = value(Xf)
    Y = value(Yf)
    F.Xnames = Xf.names
    F.Ynames = Yf.names
    return X, Y
end

getU(F::FrameSource) = value(F.Uf)
getV(F::FrameSource) = value(F.Vf)

# embed one data record
function embedU(F::FrameSource, u::Array{T,1}) where {T<:Any}
    if size(F.Uf,2) == 1
        # U has only one column, so we are embedding
        # many 1 dimensional records
        U = reshape(u, length(u),1)
    else
        # U has more than one column
        # so we must be embedding a single record
        # (and the records must have dimension equal to length(u)
        # otherwise this will fail.)
        U = reshape(u, 1, length(u))
    end
    embedU(F, U)
end

function embedU(F::FrameSource, U::Array{T,2}) where {T<:Any}
    Uf = DFrame(U, F.Uf.names)
    Xf = applyfmaplist(F.Xmaps, Uf)
    return value(Xf)
end

function embedV(F::FrameSource, U::Array{T,2}) where {T<:Any}
    Uf = DFrame(U, F.Vf.names)
    Xf = applyfmaplist(F.Ymaps, Uf)
    return value(Xf)
end
##############################################################################
# internal code

function findcolumn(c::Number, uvframe, xyframe)
    return uvframe, c
end

# by default c refers to the uvframe
# but can refer to the xyframe if
# c is only in the xyframe.names
function findcolumn(c::String, uvframe, xyframe)
    if c in uvframe.names
        return uvframe, findvalue(c, uvframe.names)
    end
    if c in xyframe.names
        return xyframe, findvalue(c, xyframe.names)
    end
    # we cannot resolve this column name
    # but return c so that it can be used in appendcol
    return xyframe, c
end

function columnvalues(DF::DFrame, col)
    j = findcolumn(col, DF)
    n = size(DF, 1)
    u = zeros(n)
    for i=1:n
        u[i] = number(DF.A[i,j])
    end
    return u
end

function appendcol(DF::DFrame, name::Int, u)
    DF.A[:,name] = u
end

# appends or replaces if the name exists already
function appendcol(DF::DFrame, name, u)
    if name != nothing && name in DF.names
        if size(u,2) > 1
            println("ERROR: Cannot specify existing destination name if using onehot")
        end
        j = findvalue(name, DF.names)
        DF.A[:,j] = u
    else
        addcolumn(DF, u)
        if size(u,2) == 1
            push!(DF.names, name)
        else
            for k=1:size(u,2)
                push!(DF.names, name * "_$k")
            end
        end
    end
end

# find all different values in column
# assign each a number 1...d
# return a dict mapping values to numbers
# and a list of values in order
function uniquecolumnvalues(DF::DFrame, col)
    j = findcolumn(col, DF)
    n = size(DF, 1)
    valtonum = Dict()
    vals = Any[]
    num = 1
    possiblevalues = sort(collect(Set(DF.A[:,j])))
    for val in possiblevalues
        valtonum[val] = num
        push!(vals, val)
        num += 1
    end
    return valtonum, vals
end

# map a column according to valtonum
function columntointegers(DF::DFrame, col, valtonum)
    j = findcolumn(col, DF)
    n = size(DF,1)
    u = zeros(n)
    for i=1:n
        u[i] = valtonum[DF.A[i,j]]
    end
    return u
end

# valtonum is a dict mapping strings to 1...d
function coltoonehot(DF::DFrame, col, valtonum)
    j = findcolumn(col, DF)
    n = size(DF,1)
    d = length(valtonum)
    u = zeros(n,d)
    for i=1:n
        u[i, valtonum[DF.A[i,j]] ] = 1
    end
    return u
end

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

abstract type FeatureMap end

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

mutable struct AddColumnFmap<:FeatureMap
    src    # source column name or number
    dest   # destination name, may be nothing
end

status(FM::AddColumnFmap) = "Add column $(FM.src) as $(FM.dest)"

function applyfmap(FM::AddColumnFmap, uvframe, xyframe)
    srcframe, srccol =  findcolumn(FM.src, uvframe, xyframe)
    u  = columnvalues(srcframe, srccol)
    appendcol(xyframe, FM.dest, u)
end

# map from dest to src
function invertfmap(FM::AddColumnFmap, uvframe, xyframe)
    destcol = findcolumn(FM.dest, xyframe)
    x = columnvalues(xyframe, destcol)
    srcframe, srccol = findcolumn(FM.src, uvframe, xyframe)
    appendcol(srcframe, FM.src, x)
end

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

mutable struct OneHotFmap<:FeatureMap
    src  # source column name or number
    dest # destination name, may be nothing
    standardize  # true if standardized
    used # if been standardized before
    mean
    std
    valtonum  # dict mapping values to numbers
    vals  # list of the values in order (i.e. a list mapping numbers to values)
end

status(FM::OneHotFmap) = "Add one-hot column $(FM.src)"

OneHotFmap(c,n,s) = OneHotFmap(c,n,s,false,nothing,nothing,nothing,nothing)

# one hot does it's own standardization since at the moment
# each feature added only takes one input column
# we do not know at the time the feature is added to the list
# how many possible values its input column will take
# and so we don't know how many output columns will be generated
# and so we cannot generate standardizers for each of them
#

function applyfmap(FM::OneHotFmap, uvframe, xyframe)
    srcframe, srccol =  findcolumn(FM.src, uvframe, xyframe)
    valtonum, vals = uniquecolumnvalues(srcframe, srccol)
    FM.valtonum = valtonum
    FM.vals = vals
    u = coltoonehot(srcframe, srccol, valtonum)
    if FM.standardize
        if !FM.used
            FM.mean = Statistics.mean(u,dims=1)
            FM.std = Statistics.std(u,dims=1)
            FM.used = true
        end
        n = size(u,1)
        u = (u - repeat(FM.mean, n, 1))./repeat(FM.std,n,1)
    end
    appendcol(xyframe, FM.dest, u)
end

# map from dest to src
function invertfmap(FM::OneHotFmap, uvframe, xyframe)
    # K is number of classes
    K = length(FM.vals)
    n = size(xyframe, 1)
    x = Array{Any}(n)
    destcols = zeros(Int64, K)
    for k=1:K
        destcols[k] = findcolumn(FM.dest * "_$(k)", xyframe)
    end
    xi = zeros(K)
    for i=1:n
        for k=1:K
            xi[k] = xyframe.A[i,destcols[k]]
        end
        maxval, maxind = findmax(xi)
        x[i] = FM.vals[maxind]
    end
    srcframe, srccol = findcolumn(FM.src, uvframe, xyframe)
    appendcol(srcframe, FM.src, x)
end

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

mutable struct OrdinalFmap<:FeatureMap
    src   # source column name or number
    dest  # destination name, may be nothing
    categories
end

status(FM::OrdinalFmap) = "Add ordinal column $(FM.src) with categories $(FM.categories)"

function applyfmap(FM::OrdinalFmap, uvframe, xyframe)
    srcframe, srccol =  findcolumn(FM.src, uvframe, xyframe)
    valtonum = Dict()
    for i=1:length(FM.categories)
        valtonum[FM.categories[i]] = i
    end
    u = columntointegers(srcframe, srccol, valtonum)
    appendcol(xyframe, FM.dest, u)
end

# maps dest to src
# map back to categories by nearest neighbor
function invertfmap(FM::OrdinalFmap, uvframe, xyframe)
    destcol = findcolumn(FM.dest, xyframe)
    x = columnvalues(xyframe, destcol)
    n = length(x)
    u = zeros(n)
    for i=1:n
        closestcat = 0
        mindist = Inf
        for j=1:length(FM.categories)
            dist = abs(x[i]-j)
            if dist < mindist
                mindist = dist
                closestcat = j
            end
        end
        u[i]  = FM.categories[closestcat]
    end
    srcframe, srccol = findcolumn(FM.src, uvframe, xyframe)
    appendcol(srcframe, FM.src, u)
end

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

mutable struct FunctionFmap<:FeatureMap
    src   # source column name or number
    dest  # destination name, may be nothing
    f
    finv
end

status(FM::FunctionFmap) = "Add function column $(FM.src)"

function FunctionFmap(src, dest, f)
    if f == log
        finv = exp
    else
        finv = nothing
    end
    return FunctionFmap(src, dest, f, finv)
end

function applyfmap(FM::FunctionFmap, uvframe, xyframe)
    srcframe, srccol =  findcolumn(FM.src, uvframe, xyframe)
    u  = columnvalues(srcframe, srccol)
    unew = [ FM.f(x) for x in u]
    appendcol(xyframe, FM.dest, unew)
end

# maps dest to src
function invertfmap(FM::FunctionFmap, uvframe, xyframe)
    destcol = findcolumn(FM.dest, xyframe)
    x = columnvalues(xyframe, destcol)
    u = [ FM.finv(a) for a in x ]
    srcframe, srccol = findcolumn(FM.src, uvframe, xyframe)
    appendcol(srcframe, FM.src, u)
end

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

mutable struct FunctionListFmap<:FeatureMap
    src  # list of source column names or numbers
    dest # destination name, may be nothing
    f
end

status(FM::FunctionListFmap) = "add function to columns: $(FM.src)"

function applyfmap(FM::FunctionListFmap, uvframe, xyframe)
    args = Any[]
    for srci in FM.src
        srcframe, srccol =  findcolumn(srci, uvframe, xyframe)
        u  = columnvalues(srcframe, srccol)
        push!(args, u)
    end
    n = length(args[1])
    unew = zeros(n)
    for i=1:n
        a2 = Any[]
        for j=1:length(args)
            push!(a2, args[j][i])
        end
        unew[i]  = FM.f(a2...)
    end
    appendcol(xyframe, FM.dest, unew)
end

########################################################
mutable struct UfunctionFmap<:FeatureMap
    dest # destination name, may be nothing
    f
end

status(FM::UfunctionFmap) = "apply function to U"

function applyfmap(FM::UfunctionFmap, uvframe, xyframe)
    n = size(uvframe, 1)
    unew  = FM.f(uvframe.A)
    appendcol(xyframe, FM.dest, unew)
end

########################################################
mutable struct AllFmap<:FeatureMap
    dest # destination name, may be nothing
    stand  # boolean
    addones #bool
    mean
    std
    used
end

AllFmap(dest, stand, addones) = AllFmap(dest, stand, addones,
                                        nothing, nothing, false)

status(FM::AllFmap) = "embed all at once"

function applyfmap(FM::AllFmap, uvframe, xyframe)
    U = uvframe.A
    if !FM.used
        if FM.stand
            FM.mean = Statistics.mean(U,dims=1)
            FM.std =Statistics.std(U,dims=1)
        else
            FM.mean = zeros(1,d)
            FM.std = ones(1,d)
        end
        FM.used = true
    end
    n,d = size(U)
    unew = (U - repeat(FM.mean,n,1))./repeat(FM.std,n,1)
    if FM.addones
        unew = [ones(n) unew]
    end
    appendcol(xyframe, FM.dest, unew)
end

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

mutable struct StandardizeFmap<:FeatureMap
    src  # source column name or number
    dest # destination name, may be nothing
    mean
    std
    used
end

status(FM::StandardizeFmap) = "standardize column $(FM.src)"

function applyfmap(FM::StandardizeFmap, uvframe, xyframe)
    srcframe, srccol =  findcolumn(FM.src, uvframe, xyframe)
    u  = columnvalues(srcframe, srccol)
    if !FM.used
        FM.mean = Statistics.mean(u)
        FM.std = Statistics.std(u)
        if FM.std<1e-2
            FM.std=1
        end
        FM.used = true
    end
    unew = (u .- FM.mean)/FM.std
    appendcol(xyframe, FM.dest, unew)
end

# map dest to src
function invertfmap(FM::StandardizeFmap, uvframe, xyframe)
    destcol = findcolumn(FM.dest, xyframe)
    x = columnvalues(xyframe, destcol)
    u = FM.std*x + FM.mean
    srcframe, srccol = findcolumn(FM.src, uvframe, xyframe)
    appendcol(srcframe, FM.src, u)
end

function StandardizeFmap(src, dest)
    return StandardizeFmap(src, dest, nothing, nothing, false)
end

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

mutable struct OneFmap<:FeatureMap
    dest # destination name, may be nothing
end

status(FM::OneFmap) = "add column of ones as $(FM.dest)"

function applyfmap(FM::OneFmap, uvframe, xyframe)
    n = size(uvframe, 1)
    u = ones(n)
    appendcol(xyframe, FM.dest, u)
end

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

function getfeature(col; name = nothing, etype="number",
                    categories = nothing,
                    f = nothing, stand=true, addones=false, kwargs...)
    if etype == "number"
        return AddColumnFmap(col, name)
    elseif etype == "onehot"
        return OneHotFmap(col, name, false)
    elseif etype == "onehotstd"
        return OneHotFmap(col, name, true)
    elseif etype == "ordinal"
        return OrdinalFmap(col, name, categories)
    elseif etype == "function"
        return FunctionFmap(col, name, f)
    elseif etype == "standardize"
        return StandardizeFmap(col, name)
    elseif etype == "product"
        return FunctionListFmap(col, name, (x,y)-> x*y)
    elseif etype == "functionlist"
        return FunctionListFmap(col, name, f)
    elseif etype == "ufunction"
        return UfunctionFmap(name, f)
    elseif etype == "all"
        return AllFmap(name, stand, addones)
    elseif etype == "one"
        return OneFmap(name)
    end
end

function addfeaturex(fmaps, count, col;
                     stand=true, name=nothing, etype="number", kwargs...)
    # make sure every feature has a name
    if name == nothing
        name = "feature$(count[1])"
        count[1] += 1
    end
    fe = getfeature(col; name=name, etype=etype, kwargs...)
    push!(fmaps, fe)
    if stand && etype != "onehot" && etype != "one" && etype != "onehotstd" && etype != "ufunction" && etype != "all"
        # add a featurizer that replaces the feature with
        # with a standardized version that has the same name
        addfeaturex(fmaps, count, name;
                      etype="standardize", stand=false, name=name, kwargs...)
    end

end

function addfeatureU(F::FrameSource, col; kwargs...)
    addfeaturex(F.Xmaps, F.count, col; kwargs...)
end

function addfeatureV(F::FrameSource, col; kwargs...)
    addfeaturex(F.Ymaps, F.count, col; kwargs...)
end

function applyfmaplist(fmaps, uvframe; estnumcols=0, verbose=false)
    n = size(uvframe,1)
    xyframe = DFrame(n; estnumcols=estnumcols)
    ne = length(fmaps)
    for i=1:ne
        if verbose
            println("FrameSource: Applying feature map: ", status(fmaps[i]))
        end
        applyfmap(fmaps[i], uvframe, xyframe)
    end
    return xyframe
end

function unembedY(F::FrameSource, y::Array{Float64,1})
    Y = reshape(y, 1, length(y))
    unembedY(F, Y)
end

function unembedY(F::FrameSource, Y::Array{Float64,2})
    Yf = DFrame(Y, F.Ynames)
    d = size(F.Vf, 2)
    n = size(Y,1)
    Vf = DFrame(zeros(n,d), F.Vf.names)
    invertfmaplist(F.Ymaps, Yf, Vf)
    V = value(Vf)
    return V
end

function invertfmaplist(fmaps, xyframe, uvframe)
    n = size(xyframe, 1)
    ne = length(fmaps)
    for i=ne:-1:1
        invertfmap(fmaps[i], uvframe, xyframe)
    end
end


text:
module DalitzPlotAnalysis
export ClebschGordon, Wignerd, WignerDϵ, Z, λ
export change_basis, cross_basis_cosθ3, cross_basis_cosθ12, cross_basis_s3

# import Base: Math.atan2
using LinearAlgebra

import GSL: sf_coupling_3j

# """
#     atan2(y::Complex{Float64}, x::Complex{Float64})
# The function calculates `atan` for a complex argument. The standard atan(y/x)
#     is either corrected by +pi or returned.
#
#     θ=2+0.1im
#     println("sin: ", cos(θ))
#     println("rec: ", cos(atan(tan(θ))))
#     println("+pi: ", cos(π+atan(tan(θ))))
#     println("at2: ", cos(atan2(sin(θ),cos(θ))))
#     println("sin: ", sin(θ))
#     println("rec: ", sin(atan(tan(θ))))
#     println("+pi: ", sin(π+atan(tan(θ))))
#     println("at2: ", sin(atan2(sin(θ),cos(θ))))
# """
# function atan2(y::Complex{Float64}, x::Complex{Float64})
#     val = atan(y/x);
#     cosval = cos(val)
#     return (real(cosval)*real(x) > 0) ? val : val+π;
# end

"""
    λ(x,y,z)

Kaellen triangle function defined by
```
λ(x,y,z) = x^2+y^2+z^2-2*x*y-2*y*z-2*z*x
```
"""
λ(x::Number,y::Number,z::Number) = x^2+y^2+z^2-2*x*y-2*y*z-2*z*x

for tp in [:Float64, :(Complex{Float64})]
    @eval function change_basis(s1::$(tp), cosθ23::Float64, m1sq::Float64, m2sq::Float64, m3sq::Float64, s::$(tp))
            # calculate s3 in (23) frame
            s3 = m1sq + m2sq +
            # 2*(  E2 * E1 -
                2*( (s1 + m2sq - m3sq)/(2*sqrt(s1)) * (s-m1sq-s1)/(2*sqrt(s1)) -
                # |p2|*cos(theta23) * (-|p1|)
                sqrt(λ(s1, m2sq, m3sq)/(4*s1))*cosθ23 * (-sqrt(λ(s, m1sq, s1)/(4*s1))) );
            # caluclate p3 in (23) frame
            p23bu = sqrt(λ(s1, m2sq, m3sq)/(4*s1));
            p3_in23_E = (s1 + m3sq - m2sq)/sqrt(4*s1);
            p3_in23_z = -p23bu*cosθ23;
            # boost to lab frame from (23) frame
            γ1 = (s + s1 - m1sq)/sqrt(4*s*s1);
            β1 = sqrt(1.0-1.0/γ1^2);
            p3_b_z = γ1*(β1*p3_in23_E+p3_in23_z);

            p3_cms = sqrt(λ(s,m3sq,s3))/sqrt(4s)
            cosθ3 = -p3_b_z/p3_cms;

            cosθ12_n = m2sq + m3sq + 2* (s3+m2sq-m1sq)/(2*sqrt(s3)) * (s-m3sq-s3)/(2*sqrt(s3)) - s1;
            cosθ12_d = (2 * sqrt(λ(s3, m1sq, m2sq)/(4*(s3))) * sqrt(λ(s, m3sq, s3)/(4*(s3))) );
            cosθ12 = cosθ12_d ≈ 0.0+0.0im ? 2.0*rand()-1.0*one(s1) : cosθ12_n/cosθ12_d;

            return s3, cosθ3, cosθ12
    end
end

function cross_basis_cosθ3(s1, cosθ23, m1sq, m2sq, m3sq, s)

     s3 = cross_basis_s3(s1, cosθ23, m1sq, m2sq, m3sq, s);
     s2 = s + m1sq + m2sq + m3sq - s1 - s3;
     # check that there will not be the Domain Error.
     λλ = λ(s, m1sq, s1)*λ(s, m3sq, s3)
     (λλ < 0) && error("λλ < 0, check masses")
     # directly get the angle
     return (2s*(m1sq+m3sq-s2)+(s+m1sq-s1)*(s+m3sq-s3))/sqrt(λλ)
end

function cross_basis_s3(s1, cosθ23, m1sq, m2sq, m3sq, s)
     # calculate s3 in (23) frame
     return m1sq + m2sq +
     # 2*(  E2 * E1 -
     2*( (s1 + m2sq - m3sq)/(2*sqrt(s1)) * (s-m1sq-s1)/(2*sqrt(s1)) -
     # |p2|*cos(theta23) * (-|p1|)
     sqrt(λ(s1, m2sq, m3sq)/(4*s1))*cosθ23 * (-sqrt(λ(s, m1sq, s1)/(4*s1))) )
end

function cross_basis_cosθ12(s1, cosθ23, m1sq, m2sq, m3sq, s)
     s3 = m1sq + m2sq +
     # 2*(  E2 * E1 -
         2*( (s1 + m2sq - m3sq)/(2*sqrt(s1)) * (s-m1sq-s1)/(2*sqrt(s1)) -
         # |p2|*cos(theta23) * (-|p1|)
         sqrt(λ(s1, m2sq, m3sq)/(4*s1))*cosθ23 * (-sqrt(λ(s, m1sq, s1)/(4*s1))) );

     cosθ12_n = m2sq + m3sq + 2* (s3+m2sq-m1sq)/(2*sqrt(s3)) * (s-m3sq-s3)/(2*sqrt(s3)) - s1;
     cosθ12_d = (2 * sqrt(λ(s3, m1sq, m2sq)/(4*(s3))) * sqrt(λ(s, m3sq, s3)/(4*(s3))) );
     # cosθ12 = cosθ12_d ≈ 0.0+0.0im ? 2.0*rand()-1.0*one(s1) : cosθ12_n/cosθ12_d;
     # cosθ12 = 2.0*rand()-1.0*one(s1)
     (cosθ12_d ≈ zero(s)) && error("λλ == 0: check your masses")
     cosθ12 = cosθ12_n/cosθ12_d

     return cosθ12
end

# recoupling function
# two options for the type:
#   - either all arguments are Float64 or energy variables are complex
for tp in [:Float64, :(Complex{Float64})]
    @eval function change_basis(s1::$(tp), cosθ1::Float64, ϕ1::Float64, cosθ23::Float64, ϕ23::Float64,
        m1sq::Float64, m2sq::Float64, m3sq::Float64, s::$(tp))
        # calculate s3 in (23) frame
        s3 = m1sq + m2sq +
        # 2*(  E2 * E1 -
            2*( (s1 + m2sq - m3sq)/(2*sqrt(s1)) * (s-m1sq-s1)/(2*sqrt(s1)) -
            # |p2|*cos(theta23) * (-|p1|)
            sqrt(λ(s1, m2sq, m3sq)/(4*s1))*cosθ23 * (-sqrt(λ(s, m1sq, s1)/(4*s1))) );
        # caluclate p3 in (23) frame
        p23bu = sqrt(λ(s1, m2sq, m3sq)/(4*s1));
        p3_in23 = [(s1 + m3sq - m2sq)/sqrt(4*s1),
                   -p23bu*sqrt(1-cosθ23^2)*cos(ϕ23),
                   -p23bu*sqrt(1-cosθ23^2)*sin(ϕ23),
                   -p23bu*cosθ23];
        # boost to lab frame from (23) frame
        γ1 = (s + s1 - m1sq)/sqrt(4*s*s1);
        β1 = sqrt(1.0-1.0/γ1^2);
        p3_b = [γ1*(p3_in23[1]+β1*p3_in23[4]),
                p3_in23[2],
                p3_in23[3],
                γ1*(β1*p3_in23[1]+p3_in23[4])];

        ct1 = cosθ1; st1 = sqrt(1.0-cosθ1^2); cp1 = cos(ϕ1); sp1 = sin(ϕ1);
        # Rz(phi1) * Ry(theta1) * p3_boost
        p3_rot = [p3_b[1],
                  cp1*ct1* p3_b[2] + (-sp1)* p3_b[3] + cp1*st1* p3_b[4],
                  sp1*ct1* p3_b[2] +   cp1 * p3_b[3] + sp1*st1* p3_b[4],
                     -st1* p3_b[2] +   0.0 * p3_b[3] +     ct1* p3_b[4]];

        cosθ3 = -p3_rot[4]/norm(p3_b[2:end]);
        ϕ3 = atan(-p3_rot[3], -p3_rot[2])  # (p3_rot[2] != zero(p3_rot[2])) ? atan2(-p3_rot[3], -p3_rot[2]) : rand()*one(p3_rot[2]);

        cosθ12_n = m2sq + m3sq + 2* (s3+m2sq-m1sq)/(2*sqrt(s3)) * (s-m3sq-s3)/(2*sqrt(s3)) - s1;
        cosθ12_d = (2 * sqrt(λ(s3, m1sq, m2sq)/(4*(s3))) * sqrt(λ(s, m3sq, s3)/(4*(s3))) );
        cosθ12 = cosθ12_d ≈ 0.0+0.0im ? 2.0*rand()-1.0*one(s1) : cosθ12_n/cosθ12_d;

        n1 = [0.0,
             -sqrt(1-cosθ1^2)*cos(ϕ1),
             -sqrt(1-cosθ1^2)*sin(ϕ1),
             -cosθ1];
        ct3 = cosθ3; st3 = sqrt(1.0-cosθ3^2); cp3 = cos(ϕ3); sp3 = sin(ϕ3);
        # Rz(phi23) * Ry(theta23) * p3_boost
        n1_rot = [0.0,#n1[1],
                  cp3*ct3* n1[2] + sp3*ct3* n1[3] + (-st3)* n1[4],
                   (-sp3)* n1[2] +    cp3 * n1[3] +   0.0 * n1[4],
                  0.0]#cp3*st3* n1[2] + sp3*st3* n1[3] +   ct3 * n1[4]];
        # println(n1_rot[3], " ", n1_rot[2])
        ϕ12 = atan(n1_rot[3], n1_rot[2])  # (n1_rot[2] != zero(n1_rot[2])) ? atan2(n1_rot[3], n1_rot[2]) : rand()*one(n1_rot[2]);
        return (s3,cosθ3,ϕ3,cosθ12,ϕ12)
    end
end

# using BenchmarkTools
# @btime change_basis(0.5834185707670233 + 0.08828205993763047im, -0.8952547092523563, 3.141592653589793, 0.1, 1.5707963267948966,0.139^2,0.139^2,0.139^2,1.4 + 0.2im)
# @btime change_basis(0.5834185707670233 + 0.08828205993763047im, -0.8952547092523563, 3.141592653589793, 0.0,
#     1.5707963267948966,0.139^2,0.139^2,0.139^2,1.4 + 0.2im)

const logfact = [sum(log(n) for n in 1:i) for i in 1:50]
lf(i) = (i>0) ? logfact[i] :  0.0;

function Wignerd(aj, am, an, β)
    (β == zero(β)) && return (am==an) ? one(β) : zero(β)
    (β == zero(β)) && return (am==-an) ? ((aj-am)%2==1 ? -one(β) : one(β))  : zero(β)
    (β ≈ π) && return (am==-an) ? ((aj-am)%2==1 ? -one(β) : one(β)) : zero(β)

    jpm = aj+am;
    jpn = aj+an;
    jmm = aj-am;
    jmn = aj-an;
    mpn = am+an;

    # common prefactor
    pref = (lf(jpm)+lf(jmm)+lf(jpn)+lf(jmn))/2.0

    #
    sb2 = sin(β/2)
    s  = log(sb2)          # log(sin(β/2))
    c  = log(1-sb2^2)/2.0  # log(cos(β/2))

    res = zero(β)
    let k0 = max(0,mpn), sgn = ((k0+jpm)%2==1 ? -1 : 1)
        for k in k0:min(jpm,jpn)
            res += sgn*exp(pref-
                    lf(jpm-k)-lf(k)-lf(k-mpn)-lf(jpn-k)+
                    (2k-mpn)*c + (jpm+jpn-2k)*s)
            sgn *= -1;
        end
    end
    res
end

# Integer argument
# Wignerd(aj::Int64, am::Int64, an::Int64, β) = Wignerd(aj//1,am//1,an//1,β)

# WignerD
WignerD(aj, am, an, α, β, γ) = Wignerd(aj,am,an,β)*cis(-(am*α+an*γ))
# WignerDϵ
function WignerDϵ(ϵP,aj, am, an, α, β, γ)
    (am < zero(am)) && return 0.0im;
    WDMp = WignerD(aj,am,an,α,β,γ)
    WDMm = (am == 0) ? WDMp : WignerD(aj,-am,an,α, β, γ);
    factor = (2ϵP-1) * (2((aj-am)%2==0)-1);
    return (WDMp - factor * WDMm) / (am == zero(am) ? 2.0 : sqrt(2.0));
end

# # Clebsches and d-function
# function ClebschGordon(j1::Rational{Int64},m1::Rational{Int64},
#                        j2::Rational{Int64},m2::Rational{Int64},
#                         j::Rational{Int64}, m::Rational{Int64})
#     factor = convert(Int64,j1+j2-m)%2==1 ? -1 : 1
#     factor*sqrt(2*j+1)*sf_coupling_3j(convert(Int64,2*j1),convert(Int64,2*j2), convert(Int64,2*j),
#                                       convert(Int64,2*m1),convert(Int64,2*m2),-convert(Int64,2*m))
# end

function ClebschGordon(two_j1::Int64,two_m1::Int64,
                       two_j2::Int64,two_m2::Int64,
                        two_j::Int64, two_m::Int64)
    factor = ((two_j1÷2)+(two_j2÷2)-(two_m÷2)) % 2 == 0 ? 1.0 : -1.0;
    factor*sqrt(two_j+1)*sf_coupling_3j(two_j1,two_j2,two_j,two_m1,two_m2,-two_m)
end

# function Wignerd(j,m1,m2,cosθ)
#     M1=m1; M2=m2; factor=1;
#     if (m1==0 && m2!=0)
#         factor *= (abs(m2)%2==1 ? -1 : 1)
#         M1=m2;M2=m1;
#     end
#     factor *= (M1 < 0) ? (abs(M1)%2==1 ? -1 : 1) : 1
#     (M2 == 0) && return factor*sqrt(4*π/(2*j+1))*sf_legendre_sphPlm(j, abs(M1), cosθ)
# end
# function Wignerd(j,m1,m2,cosθ,ϕ)
#     Wignerd(j,m1,m2,cosθ)*cis(-m1*ϕ)
# end

function Z(J,M,L,l,cosθ1,ϕ1,cosθ23,ϕ23)
    rng_lm = min(l,J);
    θ1 = acos(cosθ1)
    θ23 = acos(cosθ23)
    res = sum(ClebschGordon(2*L,0,2*l,2*lm,2*J,2*lm)*sqrt((2*L+1)*(2*l+1))*
        WignerD(J,M,lm,ϕ1,θ1,0.0)*WignerD(l,lm,0,ϕ23,θ23,0.0) for lm=-rng_lm:1:rng_lm)
    return conj(res)
end

# function Z(J::Int64,M::Int64,ϵP::Bool,L::Int64,l::Int64,cosθ1,ϕ1,cosθ23,ϕ23)
#     rng_lm = min(l,J);
#     θ1 = acos(cosθ1)
#     θ23 = acos(cosθ23)
#     res = sum(ClebschGordon(2*L,0,2*l,2*lm,2*J,2*lm)*sqrt((2*L+1)*(2*l+1))*
#         WignerDϵ(ϵP,J,M,lm,ϕ1,θ1,0.0)*WignerD(l,lm,0,ϕ23,θ23,0.0) for lm=-rng_lm:1:rng_lm)
#     return conj(res)
# end
function Z(J,M,ϵP,L,l,cosθ1,ϕ1,cosθ23,ϕ23)
    rng_lm = min(l,J);
    θ1 = acos(cosθ1)
    θ23 = acos(cosθ23)
    res = sum(ClebschGordon(2*L,0,2*l,2*lm,2*J,2*lm)*sqrt((2*L+1)*(2*l+1))*
        WignerDϵ(ϵP,J,M,lm,ϕ1,θ1,0.0)*WignerD(l,lm,0,ϕ23,θ23,0.0) for lm=-rng_lm:1:rng_lm)
    return conj(res)
end

# Z(1,0,false,0,1,0.1,0.0,0.1,0.0)

end


text:
using NeXLSpectrum
using DataFrames
using Test

@testset "MultiDet" begin
    al2o3 = loadmultispec(joinpath(@__DIR__,"MultiSpec"),"IIIE Al2O3[0]")
    sio2 = loadmultispec(joinpath(@__DIR__,"MultiSpec"),"IIIE SiO2[0]")
    fe = loadmultispec(joinpath(@__DIR__,"MultiSpec"),"IIIE Fe[0]")
    mgo = loadmultispec(joinpath(@__DIR__,"MultiSpec"),"IIIE MgO[0]")
    caf2 = loadmultispec(joinpath(@__DIR__,"MultiSpec"),"IIIE CaF2[0]")
    k412 = loadmultispec(joinpath(@__DIR__,"MultiSpec"),"IIIE K412[0]")

    comparesum(specs, sumspec) = isequal(multisum(specs).counts, sumspec.counts)
    @test comparesum(al2o3, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE Al2O3[0][4].msa")))
    @test comparesum(sio2, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE SiO2[0][4].msa")))
    @test comparesum(fe, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE Fe[0][4].msa")))
    @test comparesum(mgo, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE MgO[0][4].msa")))
    @test comparesum(caf2, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE CaF2[0][4].msa")))
    @test comparesum(k412, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE K412[0][4].msa")))

    @test isapprox(multisum(al2o3)[:LiveTime], loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE Al2O3[0][4].msa"))[:LiveTime], atol=0.001)
    @test isapprox(multisum(sio2)[:LiveTime], loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE SiO2[0][4].msa"))[:LiveTime], atol=0.001)
    @test isapprox(multisum(fe)[:LiveTime], loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE Fe[0][4].msa"))[:LiveTime], atol=0.001)
    @test isapprox(multisum(k412)[:LiveTime], loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE K412[0][4].msa"))[:LiveTime], atol=0.001)

    @test isapprox(dose(multisum(al2o3)), dose(loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE Al2O3[0][4].msa"))), atol=0.001)
    @test isapprox(dose(multisum(sio2)), dose(loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE SiO2[0][4].msa"))), atol=0.001)
    @test isapprox(dose(multisum(fe)), dose(loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE Fe[0][4].msa"))), atol=0.001)
    @test isapprox(dose(multisum(k412)), dose(loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE K412[0][4].msa"))), atol=0.001)

    comparemean(specs, sumspec) = isequal(multimean(specs).counts, sumspec.counts ./ length(specs))
    @test comparemean(al2o3, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE Al2O3[0][4].msa")))
    @test comparemean(sio2, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE SiO2[0][4].msa")))
    @test comparemean(fe, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE Fe[0][4].msa")))
    @test comparemean(mgo, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE MgO[0][4].msa")))
    @test comparemean(caf2, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE CaF2[0][4].msa")))
    @test comparemean(k412, loadspectrum(joinpath(@__DIR__,"MultiSpec","IIIE K412[0][4].msa")))

    refs = references([
        reference(n"Si", multisum(sio2), mat"SiO2"),
        reference(n"Al", multisum(al2o3), mat"Al2O3"),
        reference(n"O", multisum(al2o3), mat"Al2O3"),
        reference(n"Fe", multisum(fe), mat"Fe"),
        reference(n"Mg", multisum(mgo), mat"MgO"),
        reference(n"Ca", multisum(caf2), mat"CaF2"),
    ], 132.0 )

    q = quantify(multisum(k412), refs)
    @test isapprox(value(q.comp[n"Al"]), 0.049, atol = 0.001)
    @test isapprox(value(q.comp[n"Ca"]), 0.108, atol = 0.001)
    @test isapprox(value(q.comp[n"Fe"]), 0.080, atol = 0.001)
    @test isapprox(value(q.comp[n"Mg"]), 0.117, atol = 0.001)
    @test isapprox(value(q.comp[n"Si"]), 0.206, atol = 0.001)
    @test isapprox(value(q.comp[n"O"]),  0.459, atol = 0.001)

    @test all(isapprox(ms,v, atol=0.0001) for (ms, v) in zip(multiscore(al2o3), [0.0321, -0.0689, -0.0149, 0.0440 ]))
    @test all(isapprox(ms,v, atol=0.0001) for (ms, v) in zip(multiscore(sio2), [ 0.0478, -0.0473, -0.0287, 0.0196 ]))

    @test isapprox(multirank(k412), 0.0898, atol=0.0001)
    @test isapprox(multirank(fe), 0.0396, atol=0.0001)

end

max_stars_repo_name: NicholasWMRitchie/NeXLSpectrum

max_issues_repo_name: NicholasWMRitchie/NeXLSpectrum

max_forks_repo_name: NicholasWMRitchie/NeXLSpectrum


text:
# Generic instrument struct
abstract type Instrument end

mutable struct INSTR{instr_name} <: Instrument
    name::Symbol
    address::String
    handle# this needs to be anything, to accept PXI instruments
	initialized::Bool
	bufSize::UInt32
end
# Generic instrument constructor
INSTR(instr_name, address) = INSTR{instr_name}(instr_name, address, 0, false,UInt32(1024))

##############################################
# PSU
E364x1 = Union{INSTR{:KeysightE3645A},
	INSTR{:AgilentE3645A}
	}
E364x2 = Union{INSTR{:KeysightE3646A},
	INSTR{:AgilentE3646A},
	INSTR{:KeysightE3649A},
	INSTR{:AgilentE3649A}
	}

##############################################
# AWG
F332x1 = Union{INSTR{:Agilent33250A}
	}
F335x2 = Union{INSTR{:Agilent33500B},
	INSTR{:Keysight33500B},
	INSTR{:Agilent33522A}
	}
NIFgen = Union{INSTR{:NI5421},
	INSTR{:NI5422}
	}

##############################################
# DMM
KE200x = Union{INSTR{:Keithley2000},
	INSTR{:KE2000},
	INSTR{:Keithley2001},
	INSTR{:KE2001},
	}
NIDmm = Union{INSTR{:NI4071}
	}

##############################################
# Lecroy SCOPE
LCR6xxx = Union{INSTR{:HDO6054A},
	INSTR{:HRO64zi}
	}
# NI PXI SCOPE
NIScope = Union{INSTR{:NI5122},
			INSTR{:NI5114},
			INSTR{:NI5152},
			}


text:
module LeapSeconds

using Dates: DateTime, datetime2julian

export offset_tai_utc

include(joinpath("..", "gen", "leap_seconds.jl"))

const MJD_EPOCH = 2400000.5

# Constants for calculating the offset between TAI and UTC for
# dates between 1960-01-01 and 1972-01-01
# See ftp://maia.usno.navy.mil/ser7/tai-utc.dat

const EPOCHS = [
    36934,
    37300,
    37512,
    37665,
    38334,
    38395,
    38486,
    38639,
    38761,
    38820,
    38942,
    39004,
    39126,
    39887,
]

const OFFSETS = [
    1.417818,
    1.422818,
    1.372818,
    1.845858,
    1.945858,
    3.240130,
    3.340130,
    3.440130,
    3.540130,
    3.640130,
    3.740130,
    3.840130,
    4.313170,
    4.213170,
]

const DRIFT_EPOCHS = [
    37300,
    37300,
    37300,
    37665,
    37665,
    38761,
    38761,
    38761,
    38761,
    38761,
    38761,
    38761,
    39126,
    39126,
]

const DRIFT_RATES = [
    0.0012960,
    0.0012960,
    0.0012960,
    0.0011232,
    0.0011232,
    0.0012960,
    0.0012960,
    0.0012960,
    0.0012960,
    0.0012960,
    0.0012960,
    0.0012960,
    0.0025920,
    0.0025920,
]

"""
    offset_tai_utc(jd)

Returns the offset between International Atomic Time (TAI) and Coordinated
Universal Time (UTC) for a given Julian Date `jd`. For dates after
1972-01-01, this is the number of leap seconds.
"""
function offset_tai_utc(jd)
    mjd = jd - MJD_EPOCH

    # Before 1960-01-01
    if mjd < 36934.0
        @warn "UTC is not defined for dates before 1960-01-01."
        return 0.0
    end

    # Before 1972-01-01
    if mjd < LS_EPOCHS[1]
        idx = searchsortedlast(EPOCHS, floor(Int, mjd))
        return OFFSETS[idx] + (mjd - DRIFT_EPOCHS[idx]) * DRIFT_RATES[idx]
    end

    LEAP_SECONDS[searchsortedlast(LS_EPOCHS, floor(Int, mjd))]
end

"""
    offset_tai_utc(dt::DateTime)

Returns the offset between International Atomic Time (TAI) and Coordinated
Universal Time (UTC) for a given `DateTime`. For dates after
1972-01-01, this is the number of leap seconds.
"""
offset_tai_utc(dt::DateTime) = offset_tai_utc(datetime2julian(dt))

end


text:
include("../src/Transit.jl")
module TestTransit

using Base.Test
using DataStructures

import JSON
import Transit
import Transit.TSymbol
import Transit.TSet
import Transit.TaggedValue

function range_centered_on(n)
  range_centered_on(n, 5)
end

function range_centered_on(n, m)
  (n-m):(m+n)
end

function array_of_symbols(n, m)
  syms = map(x -> Symbol(@sprintf("key%04d", x)), 0:(n-1))
  collect(take(cycle(syms), m))
end

function dict_of_size(n)
  m = Dict{Any,Any}()
  for i in 0:(n-1)
      m[Symbol(@sprintf("key%04d", i))] = i
  end
  m
end

function list_of(a)
  result = list()
  for item in reverse(a)
    result = cons(item, result)
  end
  result
end

function dict_of(a...)
  keys = a[1:2:length(a)]
  vals = a[2:2:length(a)]

  result = Dict{Any,Any}()

  for i in 1:length(keys)
    result[keys[i]] = vals[i]
  end

  result
end

function nested_dict(n)
  Dict{Any,Any}(:f => dict_of_size(n), :s => dict_of_size(n))
end

function nested_dict_exemplar(n)
    Exemplar("map_$(n)_nested", "Map of two nested $n maps", nested_dict(n))
end

map_simple = Dict{Any,Any}(:a => 1, :b => 2, :c => 3)

map_mixed = Dict{Any,Any}(:a => 1, :b => "a string", :c => true)

map_nested = Dict{Any,Any}(:simple => map_simple, :mixed => map_mixed)

vector_simple  = Any[1, 2, 3]

vector_mixed  = Any[0, 1, 2.0, true, false, "five", :six, TSymbol(:seven), "~eight", nothing]

vector_nested = Any[vector_simple, vector_mixed]

small_strings  = ["", "a", "ab", "abc", "abcd", "abcde", "abcdef"]

powers_two =
[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384,
   32768, 65536, 131072, 262144, 524288, 1048576, 2097152, 4194304,
   8388608, 16777216, 33554432, 67108864, 134217728, 268435456,
   536870912, 1073741824, 2147483648, 4294967296, 8589934592,
   17179869184, 34359738368, 68719476736, 137438953472,
   274877906944, 549755813888, 1099511627776, 2199023255552,
   4398046511104, 8796093022208, 17592186044416, 35184372088832,
   70368744177664, 140737488355328, 281474976710656,
   562949953421312, 1125899906842624, 2251799813685248,
   4503599627370496, 9007199254740992, 18014398509481984,
   36028797018963968, 72057594037927936, 144115188075855872,
   288230376151711744, 576460752303423488, 1152921504606846976,
   2305843009213693952, 4611686018427387904, 9223372036854775808,
   18446744073709551616, 36893488147419103232]

interesting_ints = foldl(vcat, [], map(x-> Array(range_centered_on(x, 2)), powers_two))

uuids = [Base.Random.UUID( "5a2cbea3-e8c6-428b-b525-21239370dd55"),
         Base.Random.UUID( "d1dc64fa-da79-444b-9fa4-d4412f427289"),
         Base.Random.UUID( "501a978e-3a3e-4060-b3be-1cf2bd4b1a38"),
         Base.Random.UUID( "b3ba141a-a776-48e4-9fae-a28ea8571f58")]

uris = [Transit.TURI("http://example.com"),
        Transit.TURI("ftp://example.com"),
        Transit.TURI("file:///path/to/file.txt"),
        Transit.TURI("http://www.詹姆斯.com/")
]

keywords = [:a, :ab, :abc, :abcd, :abcde, :a1, :b2, :c3, :a_b]

symbols = map(x->TSymbol(x), keywords)

dates =  map(x-> Dates.unix2datetime(x), [-6106017600, 0, 946728000, 1396909037])

immutable Exemplar
  file_name::AbstractString
  description::AbstractString
  value::Any
end

exemplars = [
    Exemplar( "nil", "The nil/null/ain't there value", nothing),
    Exemplar( "true", "True", true),
    Exemplar( "false", "False", false),
    Exemplar( "zero", "Zero (integer)", 0),
    Exemplar( "one", "One (integer)", 1),
    Exemplar( "one_string", "A single string", "hello"),
    Exemplar( "one_keyword", "A single keyword", :hello),
    Exemplar( "one_symbol", "A single symbol", TSymbol(:hello)),
    Exemplar( "one_date", "A single date", Dates.unix2datetime(946728000)),
    Exemplar( "vector_simple", "A simple vector", vector_simple),
    Exemplar( "vector_empty", "An empty vector", []),
    Exemplar( "vector_mixed", "A ten element vector with mixed values", vector_mixed),
    Exemplar( "vector_nested", "Two vectors nested inside of an outter vector", vector_nested),

    Exemplar( "small_strings", "A vector of small strings", small_strings),

    Exemplar( "strings_tilde", "A vector of strings starting with ~",
	      map(x -> "~$x", small_strings)),

    Exemplar( "strings_hash", "A vector of strings starting with #",
	      map(x -> "#$x", small_strings)),

    Exemplar( "strings_hat", "A vector of strings starting with ^",
	      map(x -> "^$x", small_strings)),

    Exemplar( "small_ints", "A vector of eleven small integers", Array(range_centered_on(0))),

    Exemplar( "ints", "vector of ints", Array(0:127)),

    Exemplar( "ints_interesting", "A vector of possibly interesting positive integers",
              interesting_ints),

    Exemplar("ints_interesting_neg",
	     "A vector of possibly interesting negative integers",
	     map(x-> -1 * x, interesting_ints)),

    Exemplar( "doubles_small", "A vector of eleven doubles from -5.0 to 5.0",
      map(x->Float64(x), range_centered_on(0))),

    Exemplar( "doubles_interesting", "A vector of interesting doubles",
      [-3.14159, 3.14159, 4E11, 2.998E8, 6.626E-34]),

    Exemplar( "one_uuid", "A single UUID", first(uuids)),

    Exemplar( "uuids", "A vector of uuids", uuids),

    Exemplar( "one_uri", "A single URI", first(uris)),

    Exemplar( "uris", "A vector of URIs", uris),

    Exemplar(
      "dates_interesting",
      "A vector of interesting dates: 1776-07-04, 1970-01-01, 2000-01-01, 2014-04-07",
      dates),

    Exemplar( "symbols", "A vector of symbols", symbols),
    Exemplar( "keywords", "A vector of keywords", keywords),

    Exemplar( "list_simple", "A simple list", list_of(vector_simple)),
    Exemplar( "list_empty", "An empty list", list()),
    Exemplar( "list_mixed", "A ten element list with mixed values",
	      list_of(vector_mixed)),

    Exemplar( "list_nested", "Two lists nested inside an outter list",
      list(list_of(vector_simple), list_of(vector_mixed))),

    Exemplar( "set_simple", "A simple set", TSet(vector_simple)),
    Exemplar( "set_empty", "An empty set", TSet()),
    Exemplar( "set_mixed", "A ten element set with mixed values", TSet(vector_mixed)),
    Exemplar( "set_nested", "Two sets nested inside an outter set",
      TSet([TSet(vector_simple), TSet(vector_mixed)])),

    Exemplar( "map_simple", "A simple map", map_simple),
    Exemplar( "map_mixed", "A mixed map", map_mixed),
    Exemplar( "map_nested", "A nested map", map_nested),

    Exemplar( "map_string_keys", "A map with string keys",
              dict_of("first", 1, "second", 2, "third", 3)),

    Exemplar( "map_numeric_keys", "A map with numeric keys",
	      dict_of(1, "one", 2, "two")),

    Exemplar( "map_vector_keys", "A map with vector keys",
	      dict_of(Any[1,1], "one", Any[2,2], "two")),

    Exemplar( "map_10_items", "10 item map",  dict_of_size(10)),

    nested_dict_exemplar(10),
    nested_dict_exemplar(1935),
    nested_dict_exemplar(1936),
    nested_dict_exemplar(1937),

    Exemplar(
      "maps_two_char_sym_keys",
      "Vector of maps with identical two char symbol keys",
      [dict_of(TSymbol(:aa), 1, TSymbol(:bb), 2),
       dict_of(TSymbol(:aa), 3, TSymbol(:bb), 4),
       dict_of(TSymbol(:aa), 5, TSymbol(:bb), 6)]),

    Exemplar(
      "maps_three_char_sym_keys",
      "Vector of maps with identical three char symbol keys",
      [dict_of(TSymbol(:aaa), 1, TSymbol(:bbb), 2),
       dict_of(TSymbol(:aaa), 3, TSymbol(:bbb), 4),
       dict_of(TSymbol(:aaa), 5, TSymbol(:bbb), 6)]),

    Exemplar(
      "maps_four_char_sym_keys",
      "Vector of maps with identical four char symbol keys",
      [dict_of(TSymbol(:aaaa), 1, TSymbol(:bbbb), 2),
       dict_of(TSymbol(:aaaa), 3, TSymbol(:bbbb), 4),
       dict_of(TSymbol(:aaaa), 5, TSymbol(:bbbb), 6)]),

    Exemplar(
      "maps_two_char_keyword_keys",
      "Vector of maps with identical two char keyword keys",
      [dict_of(:aa, 1, :bb, 2),
       dict_of(:aa, 3, :bb, 4),
       dict_of(:aa, 5, :bb, 6)]),

    Exemplar(
      "maps_three_char_keyword_keys",
      "Vector of maps with identical three char keyword keys",
      [dict_of(:aaa, 1, :bbb, 2),
       dict_of(:aaa, 3, :bbb, 4),
       dict_of(:aaa, 5, :bbb, 6)]),

    Exemplar(
      "maps_four_char_keyword_keys",
      "Vector of maps with identical four char keyword keys",
      [dict_of(:aaaa, 1, :bbbb, 2),
       dict_of(:aaaa, 3, :bbbb, 4),
       dict_of(:aaaa, 5, :bbbb, 6)]),

    Exemplar(
      "maps_two_char_string_keys",
      "Vector of maps with identical two char string keys",
      [dict_of("aa", 1, "bb", 2),
       dict_of("aa", 3, "bb", 4),
       dict_of("aa", 5, "bb", 6)]),

    Exemplar(
      "maps_three_char_string_keys",
      "Vector of maps with identical three char string keys",
      [dict_of("aaa", 1, "bbb", 2),
       dict_of("aaa", 3, "bbb", 4),
       dict_of("aaa", 5, "bbb", 6)]),

    Exemplar(
      "maps_four_char_string_keys",
      "Vector of maps with identical four char string keys",
      [dict_of("aaaa", 1, "bbbb", 2),
       dict_of("aaaa", 3, "bbbb", 4),
       dict_of("aaaa", 5, "bbbb", 6)]),

    Exemplar(
      "maps_unrecognized_keys",
      "Vector of maps with keys with unrecognized encodings",
      [TaggedValue("abcde", :anything),
       TaggedValue("fghij", Symbol("anything-else"))]),

    Exemplar( "map_unrecognized_vals", "Map with vals with unrecognized encodings",
              dict_of(:key, "~Unrecognized")),

    Exemplar(
      "vector_1935_keywords_repeated_twice",
      "Vector of 1935 keywords, repeated twice",
      array_of_symbols(1935,3870)),

    Exemplar(
      "vector_1936_keywords_repeated_twice",
      "Vector of 1936 keywords, repeated twice",
      array_of_symbols(1936,3872)),

    Exemplar(
      "vector_1937_keywords_repeated_twice",
      "Vector of 1937 keywords, repeated twice",
      array_of_symbols(1937,3874)),

    Exemplar(
      "vector_unrecognized_vals",
      "Vector with vals with unrecognized encodings",
      ["~Unrecognized"]),

    Exemplar( "vector_special_numbers", "Vector with special numbers", Any[NaN, Inf, -Inf]),

    Exemplar(
      "cmap_pathological",
      "cmap pathological case discovered in transit-js and transit-cljs",
      Any[Dict(Symbol("any-value") => Dict(Any["this vector makes this a cmap"] => "any value", "any string" => :victim)),
          Dict(:victim => Symbol("any-other-value"))])
]

function issame(x::Any, y::Any)
    if x == y
        true
    else
	false
    end
end

function issame(x::TSet, y::TSet)
    if length(x) != length(y)
        return false
    else
        if x != y
            for a in y
                if !in(a, x)
                    println("Failed: $a with type $(typeof(a)) not in $x")
                end
            end
            false
        else
            true
        end
    end
end

function issame(x::AbstractFloat, y::AbstractFloat)
    if isnan(x) && isnan(y)
        true
    else
        x == y
    end
end

function issame(x::Array, y::Array)
  if length(x) != length(y)
    false
  elseif length(x) == 0
    true
  else
    for i in 1:length(x)
        if !issame(x[i], y[i])
            return false
        end
    end
    true
  end
end

function findsame(x, col)
    for item in col
	if issame(x, item)
            return true
        end
    end
    false
end

function test_reading(e::Exemplar)
  path = "../transit-format/examples/0.8/simple/$(e.file_name).json"
  actual = Transit.parse(open(path))
  if !issame(e.value, actual)
      println("\nREAD: $(e.file_name): Expected:\n$(e.value)\nBut got:\n$actual")
      return false
  end
  true
end

function test_exemplars()
    failures = 0
    for e in exemplars
      failures = test_reading(e) ? failures : failures + 1
    end

    if failures > 0
        println("Number of failures: $failures")
    end
    failures == 0
end

@test test_exemplars()

end


text:
__precompile__()
module FastTransforms

using Base, ToeplitzMatrices, Compat

import Base: *
import Compat: view

export cjt, icjt, jjt, plan_cjt, plan_icjt
export leg2cheb, cheb2leg, leg2chebu, ultra2ultra, jac2jac
export gaunt

# Other module methods and constants:
#export ChebyshevJacobiPlan, jac2cheb, cheb2jac
#export sqrtpi, pochhammer, stirlingseries, stirlingremainder, Aratio, Cratio, Anαβ
#export Cnmαβ, Cnαβ, Cnmλ, Cnλ, Λ, absf, findmindices!
#export clenshawcurtis, clenshawcurtis_plan, clenshawcurtisweights
#export fejer1, fejer_plan1, fejerweights1
#export fejer2, fejer_plan2, fejerweights2
#export RecurrencePlan, forward_recurrence!, backward_recurrence

include("fftBigFloat.jl")

include("specialfunctions.jl")
include("clenshawcurtis.jl")
include("fejer.jl")
include("recurrence.jl")

abstract FastTransformPlan{D,T}

include("ChebyshevJacobiPlan.jl")
include("jac2cheb.jl")
include("cheb2jac.jl")

include("ChebyshevUltrasphericalPlan.jl")
include("ultra2cheb.jl")
include("cheb2ultra.jl")

include("cjt.jl")

include("toeplitzhankel.jl")

leg2cheb(x...)=th_leg2cheb(x...)
cheb2leg(x...)=th_cheb2leg(x...)
leg2chebu(x...)=th_leg2chebu(x...)
ultra2ultra(x...)=th_ultra2ultra(x...)
jac2jac(x...)=th_jac2jac(x...)

include("gaunt.jl")

end # module


text:
#=
molecule:
- Julia version:
- Author: Dan
- Date: 2021-06-01
=#

mutable struct Molecule

end