text:
__precompile__()

module QuaternionArgs
using Quaternions
import Base: convert, promote_rule, show, real, imag, conj, abs, abs2, inv, +, -, /, *, rand
import Quaternions: Quaternion, imagi, imagj, imagk
import LinearAlgebra: normalize
import Random: AbstractRNG, SamplerType
import Base: isinteger, isfinite, isnan, isinf, iszero, isone

export QuaternionArg, QuaternionArgF64, QuaternionArgF32, QuaternionArgF16
export amp, phase1, phase2, phase3, normalize

struct QuaternionArg{T<:AbstractFloat} <: Number
    q::T
    phi::T
    theta::T
    psi::T
end

QuaternionArg(q::AbstractFloat,phi::AbstractFloat,theta::AbstractFloat,psi::AbstractFloat) = QuaternionArg(promote(q,phi,theta,psi)...)
QuaternionArg(q::Integer,phi::Integer,theta::Integer,psi::Integer) = QuaternionArg{Float64}(promote(q,phi,theta,psi)...)
QuaternionArg(x::AbstractFloat) = QuaternionArg(x,zero(x),zero(x),zero(x))
QuaternionArg(x::Integer) = QuaternionArg(float(x))

function Quaternion(qarg::QuaternionArg{T}) where T <: AbstractFloat
    q0=qarg.q*(cos(qarg.phi)*cos(qarg.theta)*cos(qarg.psi) + sin(qarg.phi)*sin(qarg.theta)*sin(qarg.psi))
    q1=qarg.q*(sin(qarg.phi)*cos(qarg.theta)*cos(qarg.psi) - cos(qarg.phi)*sin(qarg.theta)*sin(qarg.psi))
    q2=qarg.q*(cos(qarg.phi)*sin(qarg.theta)*cos(qarg.psi) - sin(qarg.phi)*cos(qarg.theta)*sin(qarg.psi))
    q3=qarg.q*(cos(qarg.phi)*cos(qarg.theta)*sin(qarg.psi) + sin(qarg.phi)*sin(qarg.theta)*cos(qarg.psi))

    return Quaternion(q0,q1,q2,q3)
end

function QuaternionArg(x::Quaternion{T}) where T<:Real
    QuaternionArg(convert(QuaternionF64,x))
end

function QuaternionArg(x::Quaternion{T}) where T<:AbstractFloat
    q = abs(x)
    if q==zero(typeof(q))
        return QuaternionArg(zero(typeof(q)))
    end
    x = x/q

    val = 2*(x.q1*x.q2 - x.q0*x.q3)
    # error adjustment for satisfying the domain of asin
    if val > one(typeof(q))
        val = one(typeof(q))
    elseif val < -one(typeof(q))
        val = -one(typeof(q))
    end
    psi = -asin(val)/2

    if (psi != T(pi/4)) && (psi != -T(pi/4))
        phi  = argi(x*beta((conj(x))))/2
        theta= argj(alpha(conj(x))*x)/2
    else
        phi = T(0)
        theta= argj(gamma(conj(x))*x)/2
    end

    t = Quaternion(QuaternionArg(one(T),phi,theta,psi));
    if isapprox(t,-x)
        phi = phi - sign(phi)*T(pi)
    end

    return QuaternionArg(q, phi, theta, psi)
end

const QuaternionArgF64 = QuaternionArg{Float64}
const QuaternionArgF32 = QuaternionArg{Float32}
const QuaternionArgF16 = QuaternionArg{Float16}

convert(::Type{QuaternionArg}, x::Real) = QuaternionArg(x)
convert(::Type{QuaternionArg{T}}, x::Real) where T<:AbstractFloat = QuaternionArg(x)
convert(::Type{QuaternionArg{T}}, q::QuaternionArg{S}) where {T<:AbstractFloat,S<:AbstractFloat} = QuaternionArg{T}(T(q.q), T(q.phi), T(q.theta), T(q.psi))
convert(::Type{QuaternionArg{T}}, q::QuaternionArg{T}) where T<:AbstractFloat = q
convert(::Type{QuaternionArg{T}}, q::Quaternion{T}) where T<:AbstractFloat = QuaternionArg(q)
convert(::Type{Quaternion{T}}, q::QuaternionArg{T}) where T<:Real = Quaternion(q)

promote_rule(::Type{QuaternionArg{T}}, ::Type{S}) where {T<:AbstractFloat,S<:Real} = QuaternionArg{promote_type(T,S)}
promote_rule(::Type{QuaternionArg{T}}, ::Type{QuaternionArg{S}}) where {T<:AbstractFloat,S<:Real} = QuaternionArg{promote_type(T,S)}
promote_rule(::Type{QuaternionArg{T}}, ::Type{Quaternion{S}}) where {T<:AbstractFloat,S<:Real} = Quaternion{promote_type(T,S)}

quaternionArg(q,phi,theta,psi) = QuaternionArg(q,phi,theta,psi)
quaternionArg(x) = QuaternionArg(x)
quaternionArg(q::QuaternionArg) = q

function show(io::IO, z::QuaternionArg)
    pm(z) = z < 0 ? "-$(-z)" : "+$z"
    print(io, z.q, " (phi=",pm(z.phi)," theta=", pm(z.theta), " psi=", pm(z.psi), ")")
end

QuaternionArg(x::AbstractArray{T}) where T<:QuaternionArg = x
QuaternionArg(x::AbstractArray) = copy!(similar(x,typeof(quaternion(one(eltype(x))))), x)
QuaternionArg(x::AbstractArray{T}) where T<:Quaternion = map(QuaternionArg,x)
Quaternion(x::AbstractArray{T}) where T<:QuaternionArg = map(Quaternion,x)

amp(z::QuaternionArg) = z.q
phase1(z::QuaternionArg) = z.phi
phase2(z::QuaternionArg) = z.theta
phase3(z::QuaternionArg) = z.psi
normalize(z::QuaternionArg) = QuaternionArg(one(z.q),z.phi,z.theta,z.psi)

phase1(x::AbstractVector{T}) where T<:Real = zero(x)
phase2(x::AbstractVector{T}) where T<:Real = zero(x)
phase3(x::AbstractVector{T}) where T<:Real = zero(x)

for fn in (:amp,:phase1,:phase2,:phase3,:normalize)
    @eval begin
        ($fn)(A::AbstractArray) = map(($fn),A)
    end
end

conj(z::QuaternionArg) = QuaternionArg(z.q, -z.phi, -z.theta, -z.psi)
abs(z::QuaternionArg) = z.q
abs2(z::QuaternionArg) = z.q*z.q
inv(z::QuaternionArg) = QuaternionArg(inv(Quaternion(z)))

(-)(z::QuaternionArg) = QuaternionArg(-z.q, z.phi, z.theta, z.psi)
(/)(z::QuaternionArg, x::Real) = QuaternionArg(z.q/x, z.phi, z.theta, z.psi)
(+)(z::QuaternionArg, w::QuaternionArg) = QuaternionArg(Quaternion(z)+Quaternion(w))
(-)(z::QuaternionArg, w::QuaternionArg) = QuaternionArg(Quaternion(z)-Quaternion(w))
(*)(z::QuaternionArg, w::QuaternionArg) = QuaternionArg(Quaternion(z)*Quaternion(w))
(/)(z::QuaternionArg, w::QuaternionArg) = QuaternionArg(Quaternion(z)/Quaternion(w))

rand(r::AbstractRNG, ::SamplerType{QuaternionArg{T}}) where T<:AbstractFloat =
     quaternionArg(one(T),T(2.0)*T(pi)*(rand(T)-T(0.5)),one(T)*T(pi)*(rand(T)-T(0.5)),T(0.5)*T(pi)*(rand(T)-T(0.5)))

real(z::QuaternionArg)  = Quaternion(z).q0
imagi(z::QuaternionArg) = Quaternion(z).q1
imagj(z::QuaternionArg) = Quaternion(z).q2
imagk(z::QuaternionArg) = Quaternion(z).q3

argi(z::Quaternion) = atan(z.q1,z.q0)
argj(z::Quaternion) = atan(z.q2,z.q0)
argk(z::Quaternion) = atan(z.q3,z.q0)

alpha(z::Quaternion) = Quaternion(z.q0,  z.q1, -z.q2, -z.q3)
beta(z::Quaternion)  = Quaternion(z.q0, -z.q1,  z.q2, -z.q3)
gamma(z::Quaternion) = Quaternion(z.q0, -z.q1, -z.q2,  z.q3)

isreal(q::QuaternionArg) = isreal(Quaternion(q))
isinteger(q::QuaternionArg) = isinteger(Quaternion(q))
isfinite(q::QuaternionArg) = isfinite(Quaternion(q))
isnan(q::QuaternionArg) = isnan(Quaternion(q))
isinf(q::QuaternionArg) = isinf(Quaternion(q))
iszero(q::QuaternionArg) = iszero(Quaternion(q))
isone(q::QuaternionArg) = isone(Quaternion(q))

end


text:
#!/usr/bin/julia
#
# Implementation of short iterative Lanczos method.
import LinearAlgebra: eigvals, eigvecs;
using SmoothingSplines: SmoothingSpline, fit, predict;
export KrylovBasis, ConvergenceVariation, propagate;

"Krylov basis one one-dimensional `UniformGrid`."
struct KrylovBasis{T,S,G<:UniformGrid{S,1}} <: WaveMatrix{T,S,G}
  data::Matrix{T};

  "Eigenvalues of Hamiltonian representation in Krylov space."
  λ::Vector{Float64};
  "Matrix of eigenvectors for the Hamiltonian representation in Krylov space."
  v::Matrix{Float64};

  """
  Grid spacing, requirement for WaveObject implementation and to compute inner
  products.
  """
  Δ::Number;
end
"""
    KrylovBasis{T}(grid::UniformGrid, size::Int)

Construct `KryovBasis` object with space for `size` basis vectors on
`grid`.
"""
KrylovBasis{T}(grid::G, size::Number) where {T,S,G<:UniformGrid{S}} =
  KrylovBasis{T,S,G}(zeros(T,length(grid),size), zeros(size),
                     zeros(size, size), grid.Δ);

"""
    Base.adjoint(Φ::KrylovBasis)

Compute adjoint of `Φ` and return the result as a new `KrylovBasis` object.
"""
Base.adjoint(Φ::KrylovBasis{T,S,G}) where {T,S,G} =
  KrylovBasis{T,S,G}(Base.adjoint(convert(Matrix,Φ)), Φ.λ, Φ.v, Φ.Δ);

#=
function hermiteh(n::Number, x)
  (n < 0) && throw(DomainError(n, "must be non-negative"))

  H0 = one(x)
  H1 = 2x

  if n == 0
      return H0
  elseif n == 1
      return H1
  else
      # Hm = H_m(x), Hm1 = H_{m-1}(x), Hm2 = H_{m-2}(x)
      Hm, Hm1 = H1, H0

      for m = 2:n
          Hm1, Hm2 = Hm, Hm1
          Hm = 2x*Hm1 - 2(m-1)Hm2
      end

      return Hm
  end
end

function par_model(x, p)
  res = zeros(eltype(x), length(x));
  m = length(p);
  for j ∈ 2:m res .+= p[j] * x.^(m-1); end
  return res .* exp.(-(x.-p[1]).^2/2);
end

using LsqFit;
function fit_wf(grid, y, n)
  fit = Vector{LsqFit.LsqFitResult}(undef, n+1);
  minresid = 0.0;
  minresidid = 1;
  for i ∈ 1:n
    if i == 1 p0 = [0.0, 1.0];
    elseif i == 2 p0 = [0.0, 0.0, 2.0];
    elseif i == 3 p0 = [0.0, -2.0, 0.0, 4.0];
    elseif i == 4 p0 = [0.0, 0.0, -12.0, 0.0, 8.0];
    elseif i == 5 p0 = [0.0, 12.0, 0.0, -48.0, 0.0, 16.0];
    elseif i == 6 p0 = [0.0, 0.0, 120.0, 0.0, -160.0, 0.0, 32.0];
    elseif i == 7 p0 = [0.0, -120.0, 0.0, 720.0, 0.0, -480.0, 0.0, 64.0]; end
    #p0 = zeros(eltype(y), i+1);
    fit[i] = LsqFit.curve_fit(par_model, grid, y, p0)

    resid = sum(abs.(fit[i].resid));
    if (i == 1) minresid = resid;
    elseif resid < minresid
      minredid = resid;
      minresidid = i;
    end
  end
  println("Best model was with ", minresidid-1, " order polynomials");
  return par_model(grid, fit[minresidid].param);
end=#

"""
    KrylovBasis(grid::Uniform1DGrid, H::Hamiltonian, ψ::WaveFunction, size::Int;
                debug::Bool = false)

Construct `KrylovBasis` and fill with basis vectors and appropriate
eigenvalues/eigenvectors.
"""
function KrylovBasis(grid::G, H::Hamiltonian, ψi::WaveFunction{T,1,S,G},
                     size::Real; debug::Bool = false, smooth::Bool=false,
                     smoothno::Int=2, zerofilter::Bool=false) where
          {T,S,G<:UniformGrid{S,1}}
  # Non-hermitian parameters: setting both to true prevents assumption of a
  # hermitian Hamiltonian.
  calcβ = false;  # Should the β[i] be explicitly calculated.
  orthog = false; # Explicitly orthogonalise basis at each step.

  if (size%1 != 0) throw(DomainError(size, "basis size must be an integer")); end
  if (size < 0) throw(DomainError(size, "basis size must be positive")); end
  size = convert(Int,size);

  if debug println("Generating Krylov space basis..."); end

  # Work on a copy of the passed wave function.
  ψ = deepcopy(ψi);

  # Initialise empty basis which we'll now populate.
  Φ = KrylovBasis{eltype(ψ)}(grid, size);
  #Φi = KrylovBasis{eltype(ψ)}(grid, size);
  #Φf = KrylovBasis{eltype(ψ)}(grid, size);

  # Hamiltonian representation matrix elements.
  α = zeros(Complex{Float64}, size);    # <ϕ[i]|H|ϕ[i]>
  β = zeros(Complex{Float64}, size-1);  # <ϕ[i]|H|ϕ[i-1]>

  # Vector for storing explicitly calculated β for comparison with the implicit
  # calculation method.
  if (!calcβ && debug)
    β2 = zeros(Complex{Float64}, size-1); # <ϕ[i-1]|H|ϕ[i]>
  end

  # Normalise initial wave function if not already normalised.
  nm = norm(ψ);
  if (abs(1 - nm) > 1e-6)
    Φ[:,1] = ψ ./ nm;
    if (debug) println("Normalisation of initial ψ was required, was ", nm, "."); end;
  else Φ[:,1] = ψ; end

  # Generate the basis vectors and Hamiltonian representation.
  for i ∈ 1:size
    if zerofilter && i%smoothno == 0
      for j ∈ 1:length(Φ[:,i])
        if (abs(Φ[j,i]) < 1e-14 * unit(eltype(ψ))) Φ[j,i] = zero(eltype(Φ)); continue; end
        if (abs(real(Φ[j,i])) < 1e-14 * unit(eltype(ψ))) Φ[j,i] = imag(Φ[j,i])*im; end
        if (abs(imag(Φ[j,i])) < 1e-14 * unit(eltype(ψ))) Φ[j,i] = real(Φ[j,i]); end
      end
    end
    # Use splines to smooth function to reduce accumulating errors.
    #Φi[:,i] = Φ[:,i];
    if smooth && i%smoothno == 0
      splre = predict(fit(SmoothingSpline, ustrip.(grid), real.(ustrip.(Φ[:,i])), 1e-3));
      splim = predict(fit(SmoothingSpline, ustrip.(grid), imag.(ustrip.(Φ[:,i])), 1e-3));
      Φ[:,i] = (splre + splim*im) * unit(eltype(ψ));
      #realfit = fit_wf(ustrip.(grid), ustrip.(real.(Φ[:,i])), 2*i+1);
      #imagfit = fit_wf(ustrip.(grid), ustrip.(imag.(Φ[:,i])), 2*i+1);
      #Φ[:,i] = (realfit + imagfit*im) * unit(eltype(ψ));
      #Φf[:,i] = Φ[:,i];
    end

    # Apply Hamiltonian to previous basis element.
    # ψ = H|ϕ[i]>
    ψ[:] = ustrip.(H * Φ[:,i])*unit(eltype(ψ));

    # α[i] = <ϕ[i]|H|ϕ[i]>
    α[i] = ip(Φ, i, ψ);

    # β2[i-1] = <ϕ[i-1]|H|ϕ[i]>
    if (i != 1 && calcβ)
      β[i-1] = ip(Φ, i-1, ψ);
    elseif (i != 1 && !calcβ && debug)
      β2[i-1] = ip(Φ, i-1, ψ);
    end

    # Stop after coefficient calculation for final run.
    if (i == size) break; end

    # Orthogonalise ψ over last two basis states to get new one.
    # Use modified Gram Schmidt for this procedure.
    if !orthog
      if (i != 1)
        Φ[:,i+1] = ψ - β[i-1] * Φ[:,i-1];
      else
        Φ[:,i+1] = ψ;
      end
      Φ[:,i+1] .-= α[i] * Φ[:,i];
    else # Orthogonalise over all previous basis states.
      orthog!(ψ, Φ, i);
      Φ[:,i+1] = ψ;
    end

    # Normalise.
    if (!calcβ)
      β[i] = norm(Φ, i+1);
      Φ[:,i+1] ./= β[i];
    else
      if (debug) println("Prenorm Φ[", i+1, "] = ", norm(Φ, i+1)); end
      Φ[:,i+1] ./= norm(Φ, i+1);
    end

    # Debugging info.
    if (debug)
      println("|Φ[", i+1, "]| = ", norm(Φ, i+1));
      for j ∈ 1:i
        println("<ϕ[", j, "]|Φ[", i+1, "]> = ", ip(Φ, j, Φ, i+1));
      end
    end
  end

  # Print β values from both calculation methods for comparison.
  if (!calcβ && debug) println("β = ", β); println("β2 = ", β2); end

  # Construct Hamiltonian matrix representation from elements.
  Hn = LinearAlgebra.SymTridiagonal(real.(α), real.(β));
  # Debugging.
  if (debug) println("Σ Im(α) = ", sum(imag.(α)), "; Σ Im(β) = ", sum(imag.(β))); end

  # Calculate eigenvalues.
  Φ.λ[:] = LinearAlgebra.eigvals(Hn);
  # Remove extreme eigenvalues.
  #filter!(e->abs(e)<500, λ);

  # Calculate respective eigenvectors.
  Φ.v[:,:] = LinearAlgebra.eigvecs(Hn, Φ.λ);

  # Print information about the norm of the eigenvectors.
  if (debug)
    for i ∈ 1:size
      println("Norm of eigenvector ", i, " is", LinearAlgebra.norm(Φ.v[i,:]));
    end
  end

  return Φ;
  #return (Φ, Φi, Φf);
end

"Simulation variations specific to the Lanczos method."
abstract type LanczosVariation <: SimVariation end;

"Convergence on basis application iterations."
struct ConvergenceVariation <: LanczosVariation
  ϵ::Real;
  maxeval::Int;
end

"""
    convert_to_krylov(ψ::WaveFunction{T,S,G}, Φ::KrylovBasis{K,S,G}) where
        {T,K,S,G}

Transform the supplied wave function into the Krylov basis.
"""
function convert_to_krylov(ψ::WaveFunction{T,1,S,G}, Φ::KrylovBasis{K,S,G}) where
        {T,K,S,G}
  ψv = Φ' * ψ;
  ψv ./ LinearAlgebra.norm(ψv);
end

"""
    estimate_error(sim::Lanczos, ψv::Vector, maxaccurate::Int)

Estimate the error in the Lanczos propagation step using the amplitude
of the basis vector components greater than `maxaccurate`.
"""
function estimate_error(sim::Lanczos, ψv::Vector, maxaccurate::Real)
  if (maxaccurate%1 != 0) throw(DomainError(maxaccurate, "must be integer")); end
  if !(maxaccurate isa Int) maxaccurate = convert(Int, maxaccurate); end
  ϵ = 0.0;
  for i ∈ maxaccurate:sim.basisno
    ϵ += abs(ψv[i]);
  end
  return ϵ;
end

#=function propagate(sim::Lanczos,
                   xgrid::ConfUniform1DGrid=Grid(10, 0.01),
                   par::SimParams=SimParams(0.05, 100, 1e-8),
                   H::SeparableHamiltonian=SeparableHamiltonian(xgrid, 0.5*xgrid.^2),
                   ψ::ConfWaveFunction=WaveFunction(xgrid, exp.(-xgrid.^2/2));
                   debug::Bool = true)

  # Generate a basis for the propagation.
  Φ = KrylovBasis(xgrid, H, ψ, sim.basisno);

end=#

"""
    propagate(sim::Lanczos,
              Φ::KrylovBasis{T,S,G},
              ψ::WaveFunction{T,1,S,G},
              par::SimParams=SimParams(0.05, 100, 1e-8);
              debug::Bool = true) where {T,S,G<:UniformGrid{S}}

Propagate `ψ` using the short iterative Lanczos method with the Krylov basis
vectors `Φ` provided.
"""
function propagate(sim::Lanczos,
                   Φ::KrylovBasis{T,S,G},
                   ψ::WaveFunction{T,1,S,G},
                   par::SimParams=SimParams(0.05, 100, 1e-8);
                   debug::Bool = true) where {T,S,G<:UniformGrid{S}}
  # Explicitly construct complex-valued placeholder object to prevent implicit
  # conversion to types which will raise an inexact error.
  ψv = Vector{Complex{Float64}}(undef, size(Φ)[2])
  ψv[:] = convert_to_krylov(ψ, Φ);

  # Run propagation for the specified number of basis steps.
  for i ∈ 1:sim.basissteps
    ψv[:] = Φ.v * exp(-im * par.Δt * LinearAlgebra.Diagonal(Φ.λ)) * Φ.v' * ψv;
  end

  return WaveFunction{T,1,S,G}(Φ * ψv, ψ.Δ);
end

"""

Propagate `ψ` requiring convergence for each propagation step.
"""
function propagate(sim::Lanczos,
                   conv::ConvergenceVariation,
                   Φ::KrylovBasis{T,S,G},
                   ψ::WaveFunction{T,1,S,G},
                   par::SimParams=SimParams(0.05, 100, 1e-8);
                   adptΔt::Union{AdaptiveΔt,Nothing}=nothing,
                   startΔt::Number=0,
                   debug::Bool = true) where {T,S,G<:UniformGrid{S}}
  # Define containers for representation of ψ in Krylov space used in
  # propagation.
  ψv = Vector{Complex{Float64}}(undef, size(Φ)[2])
  ψvt = Vector{Complex{Float64}}(undef, size(Φ)[2])
  ψvtp = Vector{Complex{Float64}}(undef, size(Φ)[2])
  ψvall = Array{Complex{Float64}, 2}(undef, sim.basissteps, size(Φ)[2])

  ψv[:] = convert_to_krylov(ψ, Φ);

  Δt = zeros(typeof(par.Δt), sim.basissteps+1);
  if (startΔt != 0) Δt[1] = startΔt;
  else Δt[1] = par.Δt; end

  # Vectors to save statistics for each propagation step in.
  eval = zeros(UInt,sim.basissteps); # Number of evaluations of each time step.
  numeigen = zeros(UInt,sim.basissteps); # Number of eigenvectors used.
  error = zeros(sim.basissteps); # Estimated error in propagation.

  if (adptΔt !== nothing) stopprop = false; end

  for i ∈ 1:sim.basissteps
    evaluate = true;
    while (evaluate && eval[i] < conv.maxeval)
      fill!(ψvt, 0.0);
      eval[i] += 1;
      for j ∈ 1:sim.basisno
        ψvt[:] += Φ.v[:,j] * exp(ustrip(-im * Δt[i] * Φ.λ[j])) * Φ.v[:,j]' * ψv;

        # Check if result has converged.
        if (j == 1 || sum(abs2.(ψvt - ψvtp)) > conv.ϵ)
          # Store wave function for convergence analysis in subsequent iterations.
          if (j != sim.basisno)
            ψvtp[:] = ψvt;
          # We did not have convergence after all basis vectors were applied.
          else
            println("Failed to converge at Δt = ", Δt[i], " halving...");
            Δt[i] /= 2;
            # Save number of basis vector used if this is the last evaluation.
            if (eval[i] == conv.maxeval) numeigen[i] = j; end
          end
        else # Propagation converged to new wave function.
          # Estimate error for propagation step.
          error[i] = estimate_error(sim, ψv, j);

          # Adapt time step if required.
          if adptΔt !== nothing
            (newΔt, evaluate, stopprop) = adapt_Δt(adptΔt, error[i], Δt[i], i, eval);
            # Set new time step to appropriate index depending on re-evaluation
            # status.
            if evaluate Δt[i] = newΔt;
            elseif (i < sim.basissteps+1) Δt[i+1] = newΔt; end
          # Stop evaluation for this propagation step if not adapting Δt.
          else
            evaluate = false;
            if (i < sim.basissteps+1) Δt[i+1] = Δt[i]; end
          end

          # Save number of eigenvalues used in propagation and end propagation
          # step.
          if (!evaluate)
            numeigen[i] = j;
            break;
          end
        end
      end
    end
    # Assign ψv to the propagated wave function.
    ψv[:] = ψvt;

    # Normalise wavefunction to reduce fluctuations due to additional basis
    # vectors used.
    ψv ./= LinearAlgebra.norm(ψv);

    # Record historical propagation data.
    ψvall[i,:] = ψv;

    # Stop propagation if adapting time step and limit was reached.
    if (adptΔt !== nothing && stopprop)
      if debug println("Maximum evaluation limit reached."); end
      break;
    end
  end

  # Print propagation statistics if debugging.
  if debug
    #println(get_stats_str(eval, "evaluations"));
    #println(get_stats_str(numeigen, "eigenvectors used"));
    #println(get_stats_str(error, "error"));
    #println(get_stats_str(Δt, "time step"));
  end

  return (WaveFunction{T,1,S,G}(Φ * ψv, ψ.Δ), (Δt, error, numeigen, eval));
end

ext: jl


text:
@static if VERSION >= v"1.6"
    # this union alias is added in https://github.com/JuliaLang/julia/pull/39557
    using SparseArrays: DenseMatrixUnion, AdjOrTransDenseMatrix, DenseInputVector, DenseInputVecOrMat
else
    const DenseMatrixUnion = Union{StridedMatrix, LowerTriangular, UnitLowerTriangular, UpperTriangular, UnitUpperTriangular, BitMatrix}
    const AdjOrTransDenseMatrix = Union{DenseMatrixUnion,Adjoint{<:Any,<:DenseMatrixUnion},Transpose{<:Any,<:DenseMatrixUnion}}
    const DenseInputVector = Union{StridedVector, BitVector}
    const DenseInputVecOrMat = Union{AdjOrTransDenseMatrix, DenseInputVector}
end


text:
### functions to compute missingness
"""
    missing_chunk!(x::BEDFile, irange)

A parallel execution kernel for calculating a subset of the missingness of the _genotype predictors_ of `x`.
"""
function missing_chunk!{T <: Float}(y::DenseVector{T}, x::BEDFile{T}, irange::UnitRange{Int})
    for i in irange
        y[i] = missing_col(x, i)
    end
    return y
end

"""
    missing_chunk!(x::BEDFile)

A convenience wrapper for `missing_chunk!(y, x, irange)` that automatically chooses local indexes for `irange`.
"""
missing_chunk!{T <: Float}(y::DenseVector{T}, x::BEDFile{T}) = missing_chunk!(y, x, localindexes(y))

"""
    missing_col(x::BEDFile, snp::Int)

Compute the missing of one `snp` column of a `BEDFile` object `x`.
"""
function missing_col{T <: Float}(
    x   :: BEDFile{T},
    col :: Int
)
    alleles = 0
    for i = 1:x.geno.n
        dosage = x.geno[i,col]
        if dosage == 1
            alleles += 1
        end
    end
    return alleles ./ x.geno.n
end

"""
    missing(x::BEDFile) -> Vector{Int}

This function calculates the missingness for each SNP of a `BEDFile` object `x`.
"""
function missing{T <: Float}(
    x :: BEDFile{T}
)
    y = SharedArray{T}((x.geno.p,), pids=procs(x)) :: SharedVector{T}
    @sync begin
        for q in procs(x)
            @async remotecall_wait(missing_chunk!, q, y, x)
        end
    end
    return y
end

### functions for minor allele counts

"""
    mac_chunk!(x::BEDFile, irange)

A parallel execution kernel for calculating a subset of the mean allele counts of `x`.
"""
function mac_chunk!{T <: Float}(y::DenseVector{Int}, x::BEDFile{T}, irange::UnitRange{Int})
    for i in irange
        y[i] = mac_col(x, i)
    end
    return y
end

"""
    mac_chunk!(x::BEDFile)

A convenience wrapper for `mac_chunk!(y, x, irange)` that automatically chooses local indexes for `irange`.
"""
mac_chunk!{T <: Float}(y::DenseVector{Int}, x::BEDFile{T}) = mac_chunk!(y, x, localindexes(y))

"""
    mac_col(x::BEDFile, snp::Int)

Compute the mean of one `snp` column of a `BEDFile` object `x`.
`mac_col` will ignore missing genotype values.
"""
function mac_col{T <: Float}(
    x   :: BEDFile{T},
    col :: Int
)
    alleles = 0
    #obs     = 0
    for i = 1:x.geno.n
        dosage = x.geno[i,col]
        if dosage != 1
            alleles += genoint[dosage + ONE8]
            #obs += 1
        end
    end
    return alleles
end

"""
    mac(x::BEDFile) -> Vector{Int}

This function calculates the *m*inor *a*llele *c*ounts for each SNP of a `BEDFile` object `x`.
`mac` is similar to `maf` since it counts minor alleles, but it does not return the ratios of observed alleles.
"""
function mac{T <: Float}(
    x :: BEDFile{T}
)
    #z = zeros(Int, x.geno.p)
    #@inbounds for i = 1:x.geno.p
    #    z[i] = mac_col(x, i)
    #end
    y = SharedArray{Int}((x.geno.p,), pids=procs(x)) :: SharedVector{Int}
    @sync begin
        for q in procs(x)
            @async remotecall_wait(mac_chunk!, q, y, x)
        end
    end
    return y
end

"""
    maf_chunk(y::DenseVector{T}, x::BEDFile, irange)

A parallel execution kernel for calculating a subset of the mean allele counts of `x`.
"""
function maf_chunk!{T <: Float}(y::DenseVector{T}, x::BEDFile{T}, irange::UnitRange{Int})
    for i in irange
        y[i] = maf_col(x, i)
    end
    return y
end

"""
    maf_chunk!(x::BEDFile)

A convenience wrapper for `maf_chunk!(q, x, irange)` that automatically chooses local indexes for `irange`.
"""
maf_chunk!{T <: Float}(y::DenseVector{T}, x::BEDFile{T}) = maf_chunk!(y, x, localindexes(y))

"""
    maf_col(x::BEDFile, snp::Int)

Compute the minor allele frequency of one `snp` column of a `BEDFile` object `x`.
`maf_col` will ignore missing genotype values.
"""
function maf_col{T <: Float}(
    x   :: BEDFile{T},
    col :: Int
)
    alleles = zero(T)
    obs     = 0
    for i = 1:x.geno.n
        dosage = x.geno[i,col]
        if dosage != 1
            alleles += genofloat[dosage + ONE8]
            obs += 1
        end
    end
    return alleles / (2*obs)
end

"""
    maf(x::BEDFile) -> Vector{Float}

This function calculates the *m*inor *a*llele *f*requency for each SNP of a `BEDFile` object `x`.
"""
function maf{T <: Float}(
    x :: BEDFile{T}
)
    #z = zeros(T, x.geno.p)
    #@inbounds for i = 1:x.geno.p
    #    z[i] = maf_col(x, i)
    #end
    y = SharedArray{T}((x.geno.p,)) :: SharedVector{T}
    @sync begin
        for q in procs(x)
            @async remotecall_wait(maf_chunk!, q, y, x)
        end
    end
    return y
end

"""
    sumabs2_snp(x, snp)

This function efficiently computes the squared L2 (Euclidean) norm of column `snp` of a `BEDFile` object `x`, i.e. sumabs2(x[:,snp]).

Arguments:

- `x` is the `BEDFile` object containing the compressed `n` x `p` design matrix.
- `snp` is the current SNP (column) to use in calculations.
"""
function sumabs2_snp{T <: Float}(x::BEDFile{T}, snp::Int)
    s = zero(T)       # accumulation variable, will eventually equal dot(y,z)
    t = zero(T)       # temp variable, output of interpret_genotype
    m = x.means[snp]
    d = x.precs[snp]

    # loop over all n individuals
    @inbounds for case = 1:x.geno.n
        t = x[case,snp]
        t = isnan(t) ? zero(T) : (t - m)*d
        s += t*t
    end
    return s
end

"""
    sumabs2_covariate(x, covariate)

This function efficiently computes the squared L2 (Euclidean) norm of a covariate of a `BEDFile` object `x`, i.e. sumabs2(x[:,covariate]).

Arguments:

- `x` is the `BEDFile` object containing the compressed `n` x `p` design matrix..
- `covariate` is the current nongenetic covariate (column) to use in calculations.
"""
function sumsq_covariate{T <: Float}(x::BEDFile{T}, covariate::Int)
    t = zero(T)
    s = zero(T)
    m = x.means[x.geno.p + covariate]
    d = x.precs[x.geno.p + covariate]

    # loop over all n individuals
    @inbounds for case = 1:x.geno.n
        t = (x.covar.x[case,covariate] - m) * d
        s += t*t
    end
    return s
end

"""
    sumabs2!(y, x)

Compute the squared L2 norm of each column of a compressed matrix `x` and save it to a vector `y`.

Arguments:

- `y` is the vector to fill with the squared norms.
- `x` is the `BEDFile` object that contains the compressed `n` x `p` design matrix from which to draw the columns.
"""
function Base.sumabs2!{T <: Float}(y::DenseVector{T}, x::BEDFile{T})
    p = size(x,2)
    @assert p == length(y) "y has $(length(y)) rows, x has $p columns"
    @inbounds for snp = 1:x.geno.p
        y[snp] = sumsq_snp(x,snp,x.means,x.precs)
    end
    @inbounds for covariate = 1:x.covar.p
        y[x.geno.p + covariate] = sumsq_covariate(x,covariate,x.means,x.precs)
    end
    return nothing
end

"""
    sumabs2(x::BEDFile [, shared=true, pids=procs()])

Compute the squared L2 norm of each column of a compressed matrix `x`.

Arguments:

- `x` is the `BEDFile` object that contains the compressed `n` x `p` design matrix from which to draw the columns.

Optional Arguments:
- `shared` is a `Bool` indicating whether or not to output a `SharedArray`. Defaults to `true`.
- `pids` is a vector of process IDs to which the output `SharedArray` will be distributed. Defaults to `procs()`. Has no effect when `shared = false`.
"""
function Base.sumabs2{T <: Float}(
    x       :: BEDFile{T};
    shared  :: Bool = true,
    pids    :: Vector{Int} = procs(x),
)
    p = size(x,2)
    y = ifelse(shared, SharedArray{T}((p,), init = S -> S[localindexes(S)] = zero(T), pids=pids), zeros(T, p))
    sumabs2!(y,x)
    return y
end

### functions for column means

"""
    mean_chunk(x::BEDFile, irange)

A parallel execution kernel for calculating a subset of the column means of `x`.
"""
function mean_chunk!{T <: Float}(x::BEDFile{T}, irange::UnitRange{Int})
    for i in irange
        x.means[i] = mean_col(x, i)
    end
    return nothing
end

"""
    mean_chunk!(x::BEDFile)

A convenience wrapper for `mean_chunk!(q, x, irange)` that automatically chooses local indexes for `irange`.
"""
mean_chunk!{T <: Float}(x::BEDFile{T}) = mean_chunk!(x, localindexes(x.means))

"""
    mean_col(x::BEDFile, snp::Int)

Compute the mean of one `snp` column of a `BEDFile` object `x`.
Note that unlike the normal Julia `mean` function, `mean_col` ignores `NaN`s.
"""
function mean_col{T <: Float}(x::BEDFile{T}, snp::Int)
    s = zero(T) # accumulation variable, will eventually equal mean(x,col) for current col
    t = zero(T) # temp variable, output of interpret_genotype
    u = zero(T) # count the number of people

    # loop over all n individuals
    @inbounds for case = 1:x.geno.n
        t = x[case,snp]

        # ensure that we do not count NaNs
        if isfinite(t)
            s += t
            u += one(T)
        end
    end

    # now divide s by u to get column mean and return
    return s /= u
end

"""
    mean!(x::BEDFile)

Compute the arithmetic means of the columns of a `BEDFile` object `x`.
Note that this function will ignore `NaN`s, unlike the normal Julia function `Base.mean`.

Arguments:

- `x` is the BEDFile object to use for computing column means.
"""
function Base.mean!{T <: Float}(x::BEDFile{T})
    # taken from advection example in Julia SharedArray documentation
    # this parallel execution structure distributes calculation of mean to all collaborating processors
    # each one computes an independent chunk of the mean vector based on localindexes()
    @sync begin
        for q in procs(x)
            @async remotecall_wait(mean_chunk!, q, x)
        end
    end
#    y[(x.geno.p+1):end] = vec(mean(x.covar.x,1))
    @inbounds for i = 1:x.covar.p
        x.means[x.geno.p + i] = mean(view(x.covar.x, :, i))
    end
    return nothing
end

### functions for precisions

"""
    prec_col(x::BEDFile, snp::Int)

Compute the precision of one `snp` column of a `BEDFile` object `x`.
"""
function prec_col{T <: Float}(x::BEDFile{T}, snp::Int)
    s = zero(T)      # accumulation variable, will eventually equal mean(x,col) for current col
    t = zero(T)      # temp variable, output of interpret_genotype
    u = zero(T)      # count the number of people
    m = x.means[snp] # column mean

    # loop over all n individuals
    @inbounds for case = 1:x.geno.n
        t = x[case,snp]

        # ensure that we do not count NaNs
        if isfinite(t)
            s += (t - m)^2
            u += one(T)
        end
    end

    # now compute the std = sqrt(s / (u - 1)))
    s = s <= zero(T) ? zero(T) : sqrt((u - one(T)) / s)
    return s :: T
end

"""
    prec_chunk!(x::BEDFile, irange)

A parallel execution kernel for calculating a subset of column precisions of `x`.
"""
function prec_chunk!{T <: Float}(x::BEDFile{T}, irange::UnitRange{Int})
    @inbounds for i in irange
        x.precs[i] = prec_col(x, i)
    end
    return nothing
end

"""
    prec_chunk!(x::BEDFile)

A convenience wrapper for `prec_chunk!(x, irange)` that automatically chooses local indexes for `irange`.
"""
prec_chunk!{T <: Float}(x::BEDFile{T}) = prec_chunk!(x, localindexes(x.precs))

"""
    prec!(x::BEDFile)

Compute the precision (inverse standard deviation) of the columns of a `BEDFile` object `x`.
The precisions are stored in `x.precs`.
Note that this function will ignore `NaN`s, unlike the normal Julia function `Base.std`.

Arguments:

- `x` is the `BEDFile` object to use for computing column standard deviations.
"""
function prec!{T <: Float}(x::BEDFile{T})
    # taken from advection example in Julia SharedArray documentation
    # this parallel execution structure distributes calculation of mean to all collaborating processors
    # each one computes an independent chunk of the mean vector based on localindexes()
    @sync begin
        for q in procs(x)
            @async remotecall_wait(prec_chunk!, q, x)
        end
    end
    @inbounds for i = 1:x.covar.p
        u = (one(T) / std(view(x.covar.x, :, i))) :: T
        x.precs[x.geno.p + i] = u
    end
    return nothing
end

"""
    dot(x::BEDFile, y, snp)

This function computes the dot product of a column from the `BEDFile` object `x` against a vector `y`.

Arguments:

- `x` is the `BEDFile` object with the compressed `n` x `p` design matrix.
- `y` is the vector on which to perform the dot product.
- `snp` is the desired SNP (column) of the decompressed matrix to use for the dot product.

Optional Arguments:

- `sy = sum(y)`. `sy` is used for (efficient) standardization purposes. Typically this is computed once in the execution of `x' * y`.
"""
function Base.dot{T <: Float}(
    x       :: BEDFile{T},
    y       :: DenseVector{T},
    snp     :: Int,
    sy      :: T = sum(y)
)
    s = zero(T)        # accumulation variable, will eventually equal dot(y,z)
    m = x.means[snp]   # mean of SNP predictor
    d = x.precs[snp]   # 1/std of SNP predictor

    if snp <= x.geno.p

        # loop over all individuals
        @inbounds for case = 1:x.geno.n
            t = x.geno[case,snp]

            # handle exceptions on t
            u = t == ONE8 ? zero(T) : int2geno(x,t)

            # accumulate dot product
            s += y[case] * u
        end
    else
        @inbounds for case = 1:x.geno.n
            s += x.covar[case,snp-x.geno.p] * y[case]
        end
    end

    # return the (normalized) dot product
    return (s - m*sy)*d
end

"""
    dot(x,y,snp,mask_n [, sy=sum(y), sminus=(sum(y[mask_n .== 0])])

Can also be called with a bitmask vector `mask_n` containins `0`s and `1`s which removes masked rows of `x` and `y` from the dot product.
An additional keyword argument `sminus` encodes the number of `y`s equal to `0`.
"""
function Base.dot{T <: Float}(
    x       :: BEDFile{T},
    y       :: DenseVector{T},
    snp     :: Int,
    mask_n  :: DenseVector{Int},
    sy      :: T = sum(y),
    sminus  :: T = sum(y[mask_n .== 0]) # need this for standardization purposes
)
    s = zero(T)        # accumulation variable, will eventually equal dot(y,z)
    m = x.means[snp]   # mean of SNP predictor
    d = x.precs[snp]   # 1/std of SNP predictor

    if snp <= x.geno.p

        # loop over all individuals
        @inbounds for case = 1:x.geno.n

            # only accumulate if case is not masked
            if mask_n[case] == 1
                t = x.geno[case,snp]

                # handle exceptions on t
                u = t == ONE8 ? zero(T) : int2geno(x,t)

                # accumulate dot product
                s += y[case] * u
            end
        end
    else
        @inbounds for case = 1:x.geno.n
            if mask_n[case] == 1
                s += x.covar.x[case,snp-x.geno.p] * y[case]
            end
        end
    end

    # return the (normalized) dot product
    return (s - (sy - sminus)*m)*d
end

function Base.dot{T <: Float}(
    x       :: BEDFile{T},
    y       :: DenseVector{T},
    snp     :: Int,
    mask_n  :: BitArray{1},
    sy      :: T = sum(y),
    sminus  :: T = sum(y[!mask_n]) # need this for standardization purposes
)
    s = zero(T)        # accumulation variable, will eventually equal dot(y,z)
    m = x.means[snp]   # mean of SNP predictor
    d = x.precs[snp]   # 1/std of SNP predictor

    if snp <= x.geno.p

        # loop over all individuals
        @inbounds for case = 1:x.geno.n

            # only accumulate if case is not masked
            if mask_n[case]
                t = x.geno[case,snp]

                # handle exceptions on t
                u = t == ONE8 ? zero(T) : int2geno(x,t)

                # accumulate dot product
                s += y[case] * u
            end
        end
    else
        @inbounds for case = 1:x.geno.n
            if mask_n[case]
                s += x.covar.x[case,snp-x.geno.p] * y[case]
            end
        end
    end

    # return the (normalized) dot product
    return (s - (sy - sminus)*m)*d
end

"""
    dott(x::BEDFile, b, case, idx::BitArray{1})

This function computes the dot product of a row from the `BEDFile` object `x` against a vector `b`.
It computes `dot(x[case,:], b)` as `dot(x'[:,case], b)` to respect memory stride.

Arguments:

- `x` is the `BEDFile` object with the compressed `n` x `p` design matrix.
- `b` is the vector on which to perform the dot product.
- `case` is the desired case (row) of the decompressed matrix to use for the dot product.
- `idx` is a `BitArray` that indexes `b`.
"""
function dott{T <: Float}(
    x    :: BEDFile{T},
    b    :: DenseVector{T},
    case :: Int,
    idx  :: BitArray{1}
)
    s = zero(T)     # accumulation variable, will eventually equal dot(y,z)
    @inbounds for snp = 1:x.geno.p

        # if current index of b is FALSE, then skip it since it does not contribute to Xb
        if idx[snp]

            # decompress genotype, this time from transposed matrix
            t = getindex_t(x.geno,x.geno.xt,snp,case)

            # handle exceptions on t
            u = t == ONE8 ? zero(T) : (int2geno(x,t) - x.means[snp]) * x.precs[snp]

            # accumulate dot product
            s += b[snp] * u
        end
    end
    @inbounds for covar = (x.geno.p+1):size(x,2)
        if idx[covar]
            s += b[covar] * (x.covar.xt[covar-x.geno.p,case] - x.means[covar]) * x.precs[covar]
        end
    end

    # return the dot product
    return s
end

"""
    A_mul_B!(xb, x, b, idx, k, mask_n)

Can also be called with a bitmask vector `mask_n` containins `0`s and `1`s which excludes or includes (respectively) elements of `x` and `b` from the dot product.
"""
function Base.A_mul_B!{T <: Float}(
    xb     :: DenseVector{T},
    x      :: BEDFile{T},
    b      :: DenseVector{T},
    idx    :: BitArray{1},
    k      :: Int,
    mask_n :: DenseVector{Int}
)
    # error checking
    @assert 0 <= k <= size(x,2) "Number of active predictors must be nonnegative and less than p"
    @assert k >= sum(idx) "Must have k >= sum(idx) or X*b will not compute correctly"
    n = length(xb)
    @assert n == x.geno.n "xb has $n rows but x has $(x.geno.n) rows"

    # loop over the desired number of predictors
    @inbounds for case = 1:n
        if mask_n[case] == 1
            xb[case] = dott(x, b, case, idx)
        end
    end

    return nothing
end

# easy wrapper for previous function
A_mul_B!{T <: Float}(xb::DenseVector{T}, x::BEDFile{T}, b::DenseVector{T}) = A_mul_B!(xb, x, b, b .!= 0, countnz(b), x.mask)

function Base.A_mul_B!{T <: Float}(
    xb     :: DenseVector{T},
    x      :: BEDFile{T},
    b      :: DenseVector{T},
    idx    :: BitArray{1},
    k      :: Int,
    mask_n :: BitArray{1}
)
    # error checking
    @assert 0 <= k <= size(x,2) "Number of active predictors must be nonnegative and less than p"
    @assert k >= sum(idx) "Must have k >= sum(idx) or X*b will not compute correctly"
    n = length(xb)
    @assert n == x.geno.n "xb has $n rows but x has $(x.geno.n) rows"

    # loop over the desired number of predictors
    @inbounds for case = 1:n
        if mask_n[case]
            xb[case] = dott(x, b, case, idx)
        end
    end

    return nothing
end

"""
    A_mul_B!(xb, x, b, idx, k)

This function computes the operation `x*b` for the compressed `n` x `p` design matrix from a `BEDFile` object.
`A_mul_B!()` respects memory stride for column-major arrays by using a compressed transpose of genotypes and an uncompressed transpose of covariates.
It also assumes a sparse b, for which we have a `BitArray` index vector `idx` to select the nonzeroes.

Arguments:

- `xb` is the `n`-dimensional output vector.
- `x` is the `BEDFile` object for the compressed `n` x `p` design matrix.
- `b` is the `p`-dimensional vector against which we multiply `x`.
- `idx` is a `BitArray` that indexes the nonzeroes in `b`.
- `k` is the number of nonzeroes to use in computing `x*b`.
"""
function Base.A_mul_B!{T <: Float}(
    xb   :: DenseVector{T},
    x    :: BEDFile{T},
    b    :: DenseVector{T},
    idx  :: BitArray{1},
    k    :: Int
)
    # error checking
    n = length(xb)
    @assert n == x.geno.n "xb has $n rows but x has $(x.geno.n)"
    @assert 0 <= k <= size(x,2) "Number of active predictors must be nonnegative and less than p"
    @assert k >= sum(idx) "Must have k <= sum(idx) or X*b will not compute correctly"

    # loop over the desired number of predictors
    @inbounds for case = 1:x.geno.n
        xb[case] = dott(x, b, case, idx)
    end

    return nothing
end

"""
    A_mul_B(x, b, idx, k, mask_n)

Can also be called with a bitmask vector `mask_n` containins `0`s and `1`s which excludes or includes (respectively) elements of `x` and `b` from the dot product.
"""
function A_mul_B{T <: Float}(
    x      :: BEDFile{T},
    b      :: SharedVector{T},
    idx    :: BitArray{1},
    k      :: Int,
    mask_n :: DenseVector{Int}
)
    xb = SharedArray{T}((x.geno.n,), init = S -> S[localindexes(S)] = zero(T), pids=procs(x)) :: SharedVector{T}
    A_mul_B!(xb, x, b, idx, k, mask_n)
    return xb
end

function A_mul_B{T <: Float}(
    x      :: BEDFile{T},
    b      :: Vector{T},
    idx    :: BitArray{1},
    k      :: Int,
    mask_n :: DenseVector{Int};
)
    xb = zeros(T, x.geno.n)
    A_mul_B!(xb, x, b, idx, k, mask_n)
    return xb
end

"""
    A_mul_B(x::BEDFile, b, idx, k)

This function computes the operation `x*b` for the compressed `n` x `p` design matrix from a `BEDFile` object.
`A_mul_B()` respects memory stride for column-major arrays.
It also assumes a sparse `b`, for which we have a `BitArray` index vector `idx` to select the nonzeroes.
The output type matches the type of `b`.

Arguments:

- `x` is the `BEDFile` object for the compressed `n` x `p` design matrix.
- `b` is the `p`-dimensional vector against which we multiply `x`.
- `idx` is a `BitArray` that indexes the nonzeroes in `b`.
- `k` is the number of nonzeroes to use in computing `x*b`.
"""
function A_mul_B{T <: Float}(
    x    :: BEDFile{T},
    b    :: SharedVector{T},
    idx  :: BitArray{1},
    k    :: Int
)
    xb = SharedArray{T}((x.geno.n,), init = S -> S[localindexes(S)] = zero(T), pids=procs(x)) :: SharedVector{T}
    A_mul_B!(xb, x, b, idx, k)
    return xb
end

function A_mul_B{T <: Float}(
    x   :: BEDFile{T},
    b   :: Vector{T},
    idx :: BitArray{1},
    k   :: Int;
)
    xb = zeros(T, x.geno.n)
    A_mul_B!(xb, x, b, idx, k)
    return Xb
end

# easy wrapper for A_mul_B!
A_mul_B{T <: Float}(x::BEDFile{T}, b::DenseVector{T}) = A_mul_B(x, b, b .!= 0, countnz(b))

# overload matrix-vector multiplication x*b with a BEDFile
#Base.*{T <: Float}(x :: BEDFile{T}, b :: DenseVector{T}) = A_mul_B(x, b, b .!= zero(T), countnz(b), pids=procs(x))

"""
    At_mul_B_chunk!(xty, x::BEDFile, y, mask_n, irange)

A parallel execution kernel for calculating `x' * y` with a `BEDFile` object `x` that allows inclusion/exclusion of indices based on `mask_n`.
"""
function At_mul_B_chunk!{T <: Float}(
    xty    :: SharedVector{T},
    x      :: BEDFile{T},
    y      :: SharedVector{T},
    mask_n :: DenseVector{Int},
    irange :: UnitRange{Int},
    sy     :: T = sum(y),
    sminus :: T = begin z = mask_n .== 0; sum(y[z]) end
)
    @inbounds for i in irange
        xty[i] = dot(x, y, i, mask_n, sy, sminus)
    end
    return nothing
end

function At_mul_B_chunk!{T <: Float}(
    xty    :: SharedVector{T},
    x      :: BEDFile{T},
    y      :: SharedVector{T},
    mask_n :: BitArray{1},
    irange :: UnitRange{Int},
    sy     :: T = sum(y),
    sminus :: T = sum(y[!mask_n])
)
    @inbounds for i in irange
        xty[i] = dot(x, y, i, mask_n, sy, sminus)
    end
    return nothing
end

"""
    At_mul_B_chunk!(xty, x::BEDFile, y, mask_n)

A convenience wrapper for `At_mul_B_chunk!(xty, x, y, irange)` that automatically chooses local indexes for `irange` and excludes elements based on `mask_n`.
"""
function At_mul_B_chunk!{T <: Float}(
    xty    :: SharedVector{T},
    x      :: BEDFile{T},
    y      :: SharedVector{T},
    mask_n :: DenseVector{Int},
    sy     :: T = sum(y),
    sminus :: T = begin z = mask_n .== 0; sum(y[z]) end
)
    At_mul_B_chunk!(xty, x, y, mask_n, localindexes(xty), sy, sminus)
    return nothing
end

function At_mul_B_chunk!{T <: Float}(
    xty    :: SharedVector{T},
    x      :: BEDFile{T},
    y      :: SharedVector{T},
    mask_n :: BitArray{1},
    sy     :: T = sum(y),
    sminus :: T = sum(y[!mask_n])
)
    At_mul_B_chunk!(xty, x, y, mask_n, localindexes(xty), sy, sminus)
    return nothing
end

"""
    At_mul_B!(xty, x::BEDFile, y, mask_n, [, pids=procs()])

Can also be called with a bitmask vector `mask_n` containins `0`s and `1`s which excludes or includes elements of `x` and `y` from the dot product.
"""
function Base.At_mul_B!{T <: Float}(
    xty     :: SharedVector{T},
    x       :: BEDFile{T},
    y       :: SharedVector{T},
    mask_n  :: DenseVector{Int};
    pids    :: Vector{Int} = procs(x),
    sy      :: T = sum(y),
    #sminus  :: T = sum(y[mask_n .== 0])
    sminus  :: T = begin z = mask_n .== 0; sum(y[z]) end
)
    # error checking
    p = size(x,2)
    @assert p <= length(xty) "Attempting to fill argument xty of length $(length(xty)) with $p elements!"
    @assert x.geno.n == length(y) "Argument y has $(length(y)) elements but should have $(x.geno.n) of them!"
    @assert pids == procs(xty) == procs(y) == procs(x.geno.x) "SharedArray arguments to At_mul_B! must be seen by same processes"

    # each processor will compute its own chunk of x'*y
    @sync begin
        for q in procs(xty)
#            @async remotecall_wait(At_mul_B_chunk!, q, xty, x, y, mask_n, sy, sminus)
            @async remotecall_wait(At_mul_B_chunk!, q, xty, x, y, mask_n, sy, sminus)
        end
    end
    return nothing
end

function Base.At_mul_B!{T <: Float}(
    xty     :: SharedVector{T},
    x       :: BEDFile{T},
    y       :: SharedVector{T},
    mask_n  :: BitArray{1};
    pids    :: Vector{Int} = procs(x),
    sy      :: T = sum(y),
    #sminus  :: T = sum(y[mask_n .== 0])
    sminus  :: T = begin z = !mask_n; sum(y[z]) end
)
    # error checking
    p = size(x,2)
    @assert p <= length(xty) "Attempting to fill argument xty of length $(length(xty)) with $p elements!"
    @assert x.geno.n == length(y) "Argument y has $(length(y)) elements but should have $(x.geno.n) of them!"
    @assert pids == procs(xty) == procs(y) == procs(x.geno.x) "SharedArray arguments to At_mul_B! must be seen by same processes"

    # each processor will compute its own chunk of x'*y
    @sync begin
        for q in procs(xty)
#            @async remotecall_wait(At_mul_B_chunk!, q, xty, x, y, mask_n, sy, sminus)
            @async remotecall_wait(At_mul_B_chunk!, q, xty, x, y, mask_n, sy, sminus)
        end
    end
    return nothing
end

function Base.At_mul_B!{T <: Float}(
    xty     :: Vector{T},
    x       :: BEDFile{T},
    y       :: Vector{T},
    mask_n  :: DenseVector{Int},
    sy      :: T = sum(y),
    sminus  :: T = begin z = mask_n .== 0; sum(y[z]) end
)
    # error checking
    p = size(x,2)
    @assert p <= length(xty) "Attempting to fill argument xty of length $(length(xty)) with $p elements!"
    @assert x.geno.n == length(y) "Argument y has $(length(y)) elements but should have $(x.geno.n) of them!"

    # loop over the desired number of predictors
    @inbounds for snp = 1:p
        xty[snp] = dot(x, y, snp, mask_n, sy, sminus)
    end
    return nothing
end

function Base.At_mul_B!{T <: Float}(
    xty     :: Vector{T},
    x       :: BEDFile{T},
    y       :: Vector{T},
    mask_n  :: BitArray{1},
    sy      :: T = sum(y),
    #sminus  :: T = sum(y[mask_n .== 0])
    sminus  :: T = begin z = !mask_n; sum(y[z]) end
)
    # error checking
    p = size(x,2)
    @assert p <= length(xty) "Attempting to fill argument xty of length $(length(xty)) with $p elements!"
    @assert x.geno.n == length(y) "Argument y has $(length(y)) elements but should have $(x.geno.n) of them!"

    # loop over the desired number of predictors
    @inbounds for snp = 1:p
        xty[snp] = dot(x, y, snp, mask_n, sy, sminus)
    end
    return nothing
end

"""
    At_mul_B_chunk!(xty, x::BEDFile, y, irange)

A parallel execution kernel for calculating `x' * y` with a `BEDFile` object `x`.
"""
function At_mul_B_chunk!{T <: Float}(
    xty     :: SharedVector{T},
    x       :: BEDFile{T},
    y       :: SharedVector{T},
    irange  :: UnitRange{Int},
    sy      :: T = sum(y),
)
    @inbounds for i in irange
        xty[i] = dot(x, y, i, sy)
    end
    return nothing
end

"""
    At_mul_B_chunk!(xty, x::BEDFile, y)

A convenience wrapper for `At_mul_B_chunk!(xty, x, y, irange)` that automatically chooses local indexes.
"""
function At_mul_B_chunk!{T <: Float}(
    xty     :: SharedVector{T},
    x       :: BEDFile{T},
    y       :: SharedVector{T},
    sy      :: T = sum(y)
)
    At_mul_B_chunk!(xty, x, y, localindexes(xty), sy)
    return nothing
end

"""
    At_mul_B!(xty, x, y, [, pids=procs()])

This function computes the operation `x'*y` for the compressed `n` x `p` design matrix from a `BEDFile` object.
`At_mul_B!` enforces a uniform type (`SharedArray` v. `Array` and `Float64` v. `Float32`) across all arrays.
It also requires `SharedArray` arguments to be seen by the same processes, i.e. `procs(xty) == procs(y) == procs(x)`.

Arguments:

- 'xty' is the p-dimensional output vector.
- `x` is the `BEDFile` object for the compressed `n` x `p` design matrix.
- `y` is the `n`-dimensional vector against which we multiply `x`.

Optional Arguments:

- `pids` is a vector of process IDs over which to distribute the `SharedArray`s for `means` and `precs`, if not supplied,
   as well as the output vector. Defaults to `procs()`. Only available for `SharedArray` arguments to `xty` and `y`.
"""
function Base.At_mul_B!{T <: Float}(
    xty     :: SharedVector{T},
    x       :: BEDFile{T},
    y       :: SharedVector{T};
    pids    :: Vector{Int} = procs(x),
    sy      :: T = sum(y)
)
    # error checking
    p = size(x,2)
    @assert p <= length(xty) "Attempting to fill argument xty of length $(length(xty)) with $p elements!"
    @assert x.geno.n == length(y) "Argument y has $(length(y)) elements but should have $(x.n) of them!"
    @assert pids == procs(xty) == procs(y) == procs(x) "SharedArray arguments to At_mul_B! must be seen by same processes"

    # each processor will compute its own chunk of x'*y
    @sync begin
        for q in procs(xty)
#            @async remotecall_wait(At_mul_B_chunk!, q, xty, x, y, sy)
            @async remotecall_wait(At_mul_B_chunk!, q, xty, x, y, sy)
        end
    end

    return nothing
end

function Base.At_mul_B!{T <: Float}(
    xty     :: Vector{T},
    x       :: BEDFile{T},
    y       :: Vector{T};
    sy      :: T   = sum(y),
)
    # error checking
    p = size(x,2)
    @assert p <= length(xty) "Attempting to fill argument xty of length $(length(xty)) with $p elements!"
    @assert x.geno.n == length(y) "Argument y has $(length(y)) elements but should have $(x.geno.n) of them!"

    # loop over the desired number of predictors
    @inbounds for snp = 1:p
        xty[snp] = dot(x, y, snp, sy)
    end
    return nothing
end

# easy function wrapper
At_mul_B!{T <: Float}(xty::DenseVector{T}, x::BEDFile{T}, y::DenseVector{T}) = At_mul_B!(xty, x, y, x.mask)

"""
    At_mul_B(x, y, mask_n [, pids=procs()])

This function computes `x'*y` for the compressed `n` x `p` design matrix from a `BEDFile` object.
It uses a bitmask `mask_n` to exclude certain rows of `x` from the calculations.

Arguments:

- `x` is the `BEDFile` object for the compressed `n` x `p` design matrix.
- `y` is the `n`-dimensional vector against which we multiply `x`.
- `mask_n` is an `n`-vector of `0`s and `1`s that indicates which rows to exclude or include, respectively.

Optional Arguments:

- `pids` is a vector of process IDs over which to distribute the `SharedArray`s for `means` and `precs`, if not supplied,
   as well as the output vector. Defaults to `procs()`. Only available for `SharedArray` arguments to `b`.
"""
function Base.At_mul_B{T <: Float}(
    x       :: BEDFile{T},
    y       :: SharedVector{T},
    mask_n  :: DenseVector{Int};
    pids    :: Vector{Int} = procs(x),
    sy      :: T = sum(y)
)
    p = size(x,2)
    xty = SharedArray{T}((p,), init = S -> S[localindexes(S)] = zero(T), pids=pids) :: SharedVector{T}
    At_mul_B!(xty, x, y, mask_n, pids=pids, sy=sy)
    return xty
end

function Base.At_mul_B{T <: Float}(
    x       :: BEDFile{T},
    y       :: Vector{T},
    mask_n  :: DenseVector{Int};
    sy      :: T = sum(y),
    sminus  :: T = sum(y[mask_n .== 0])
)
    p = size(x,2)
    xty = zeros(T, p)
    At_mul_B!(xty, x, y, mask_n, sy, sminus)
    return xty
end

At_mul_B{T <: Float}(x::BEDFile{T}, y::DenseVector{T}) = At_mul_B(x, y, x.mask)

max_stars_repo_name: klkeys/PLINK

max_issues_repo_name: klkeys/PLINK


text:
#############################################################################
# Web client messaging functions
#############################################################################

"""
The `showmsg` function converts the argument into an appropriate web component
creation message to the web client.
Default method tries to find the `MIME"text/html"` or `MIME"text/plain"` representation
and use it as the html parameter of an `html-node` web component.
Specialize this function by the `obj` argument type to create specific web components.
(see Katex type for an example)
"""
function showmsg(obj::Any, opts::Dict=Dict())
  buf = IOBuffer()
  if method_exists(show, (IO, MIME"text/html", typeof(obj)))
    show(buf, MIME"text/html"(), obj)
  elseif method_exists(show, (IO, MIME"text/plain", typeof(obj)))
    show(buf, MIME"text/plain"(), obj)
  else
    error("no show function found for type $(typeof(obj))")
  end

  args = Dict(:newnid   => getnid(),
              :compname => "html-node",
              :params   => Dict(:html => String(take!(buf))),
              :deco     => opts)

  sendcurrent("append", args)
  nothing
end

"""
The `loadmsg` function instructs the web client to load the necessary javascript
code defining a new component in the current session.
This function is called for each specified type given as argument to the
`@redirect` macro call.
Default method does nothing since the default web component `html-node` is
present by default in the Vuejs client.
Specialize this function by the argument type to load new web components in the
current session (see Katex type).
"""
function loadmsg(::Type)
end


text:
__precompile__()
module Utils

export convert_tmpfile, defaultinputdict!, printparams, initlogo

using SeisIO, Printf, Dates, JLD2, FileIO

"""
convert_tmpfile(InputDict::Dict)

convert temporal file in "./seisdownload_tmp" to prescribed format.
It has salvage mode, which allows to compile the temporal files in the case of failing during the download.
"""
function convert_tmpfile(InputDict::Dict; salvage::Bool=false)

    println("-------START CONVERTING-------------")

	# save data to fopath file
	fopath = InputDict["fopath"]

	t = jldopen(InputDict["finame"])

	file = jldopen(fopath, "w")

	file["info/stationlist"]     = t["info/stationlist"];

	if InputDict["IsStartendtime"]
		file["info/DLtimestamplist"] = InputDict["DLtimestamplist_selected"];
		file["info/starttime"]       = InputDict["starttime"];
		file["info/endtime"]         = InputDict["endtime"];
	else
		file["info/DLtimestamplist"] = t["info/DLtimestamplist"];
		file["info/starttime"]       = t["info/starttime"];
		file["info/endtime"]         = t["info/endtime"];
	end

	JLD2.close(t)

	# find all temporal files
    paths = ls(InputDict["tmppath"])
    fmt = InputDict["outputformat"]

    stationlist     = []
    DLtimestamplist = []
    varnamelist     = []

	IsIsolateComponents = InputDict["IsIsolateComponents"]
	isostationlist = []
    for path in paths
        #println(path)

		# split path
		tmpname = split(path, "/")[end]

		ftmpname = split(tmpname, ".")
		y, d,net, sta, loc, cha = split(tmpname, ".")
		iso_stationinfo = (join([y, d, net, sta, loc], "-"), cha)
		# if occursin("-", ftmpname[3])
		# 	# # format would be y, jd, T00:00:00, sta, loc, cha
		# 	y, d, tmpT, net, sta, loc, cha = split(tmpname, ".")
		# 	iso_stationinfo = (join([y, d, net, sta, loc], "-"), cha)
		# elseif !occursin("-", ftmpname[3]) && IsIsolateComponents == true
		# 	@warn "format of tmp file is not y, jd, time, sta, loc, cha. Turn off IsIsolateComponents."
		# 	IsIsolateComponents = false
		# 	iso_stationinfo = []
		# end

		#println(iso_stationinfo)

        try
            S = SeisIO.rseis(path)[1]
            for ii = 1:S.n #loop at each seis channel
				if IsIsolateComponents
					# check channel isolation
					conflictsta = filter(x -> x[1]==iso_stationinfo[1] && string(x[2][end][end])==string(iso_stationinfo[2][end]), isostationlist)
				end

				if !isempty(conflictsta)
					# here this channel has conflicting channel such as same day, same components but different channel.
					# if IsIsolateComponents == true, skip to avoid multiple stations for the purpose of cross-correlation.
					if IsIsolateComponents
						sta1 = join([iso_stationinfo[1], iso_stationinfo[2]], "-")
						sta2 = join([conflictsta[1][1], conflictsta[1][2]], "-")
						txt = @sprintf("Two identical location but different channel were found.\n%s :%s: Discard %s.",
									sta1, sta2, sta1)
						println(txt)
						#rm(path)
						continue;
					end
				end

				if IsIsolateComponents
					# update isostationlist
					push!(isostationlist, iso_stationinfo)
				end
				# println(S[ii])
                # make station list
                staid = S[ii].id

                # save data (checking whether it's already in the jld2 because it causes an error)
                #parse info
                s_str = string(u2d(S[ii].t[1,2]*1e-6))

                # select output format
                if fmt == "JLD2"
                    yj = parse(Int64, s_str[1:4])
                    mj = parse(Int64, s_str[6:7])
                    dj = parse(Int64, s_str[9:10])
                    tj = string(s_str)[11:19]

                    djm2j = md2j(yj, mj, dj)
                    # groupname = string(yj)*"."*string(djm2j)*"."*replace(tj,":"=>"-") #Year_Julianday_Starttime
					groupname = string(yj)*"."*string(djm2j)
                    varname = joinpath(groupname, staid)
					# println(groupname)
                    if isempty(filter(x -> x==varname, varnamelist))
                        push!(varnamelist, varname)
                        file[varname] = S[ii]
                    end
                    # @info "save data $varname"
                else
                    error("output format in $fmt is not implemented yet.")
                end
            end

			rm(path)

        catch y
            println(y)
        end
    end

    JLD2.close(file)

	rm(InputDict["tmppath"], recursive=true, force=true)

    return nothing
end

"""
printparams(param::Dict)

print parameters
"""
function printparams(param::Dict)
    printstyled("-----------Input Parameters-----------\n"; color=:cyan, bold=true)
    for key in keys(param)
        if length(string(param["$key"])) > 60
            param_str = string(param["$key"])[1:30]*"..."
        else
            param_str = string(param["$key"])
        end
        println(@sprintf("%-24s = %-10s", key, param_str))
    end
end

"""
defaultinputdict(InputDict::Dict)

default if input parameter is missing.
"""
function defaultinputdict!(InputDict::Dict)

	# default values
	def = Dict()
	def["finame"]				 	= "./input.jld2"
	def["IsKurtosisRemoval"] 		= true
	def["max_edgetaper_duration"] 	= 60 * 5
	def["kurtosis_tw_sparse"] 		= 60
	def["kurtosis_timewindow"] 		= 60*3
	def["kurtosis_threshold"] 		= 2.0
	def["IsSTALTARemoval"] 			= true
	def["stalta_longtimewindow"] 	= 60*60*2
	def["stalta_threshold"] 		= 1.2
	def["stalta_absoluteclip"] 		= 1e20
	def["max_wintaper_duration"] 	= 60 * 3
	def["removal_shorttimewindow"] 	= 60 * 3
	def["overlap"] 					= 60
	def["IsIsolateComponents"] 		= false
	def["IsSaveFig"] 				= false
	def["plot_kurtosis_α"] 			= 1.2
	def["plot_boxheight	"] 			= 1.5
	def["plot_span"] 				= 100
	def["outputformat"]				= "JLD2"
	def["IsStartendtime"] 			= false
	def["fodir"] 					= "./dataset"
    def["foname"] 					= "eq_removed.jld2"

	for key in keys(def)
		if !haskey(InputDict, key)
			InputDict["$key"] = def["$key"]
		end
	end

end

"""
initlogo()

print initial logo
"""
function initlogo()

    print("
    _____      _      ______
   /  ___|    (_)     | ___ \\
   \\ `--.  ___ _ ___  | |_/ /___ _ __ ___   _____   _____
    `--. \\/ _ \\ / __| |    // _ \\ '_ ` _ \\ / _ \\ \\ / / _ \\
   /\\__/ /  __/ \\__ \\ | |\\ \\  __/ | | | | | (_) \\ V /  __/
   \\____/ \\___|_|___/ \\_| \\_\\___|_| |_| |_|\\___/ \\_/ \\___|
    ______ ____
   |  ____/ __ \\       |
   | |__ | |  | |      |
   |  __|| |  | |      |  v1.0 (Last update 07/07/2019)
   | |___| |__| |      |  © Kurama Okubo
   |______\\___\\_\\      |

")

    println("Job start running at "*string(now())*"\n")

end

end


text:
using ForwardDiff, Calculus
# moments should operate on a vector and return a nXg matrix
# of moment contributions

# ordinary GMM with provided weight matrix
# weight should be gXg positive definite
function gmm(moments, θ, weight)
    # average moments
    m = θ -> vec(mean(moments(θ),dims=1)) # 1Xg
    # moment contributions
    # GMM criterion
    obj = θ -> dot(m(θ), weight, m(θ))
    # do minimization
    θhat, objvalue, converged = fminunc(obj, θ)
    # derivative of average moments
    D = try
        ForwardDiff.jacobian(m, vec(θhat))'
    catch
        Calculus.jacobian(m, vec(θhat), :central)'
    end
    # moment contributions at estimate
    ms = moments(θhat)
    return θhat, objvalue, D, ms, converged
end

# CUE GMM, weight computed by NW
function gmm(moments, θ)
    # average moments
    m = θ -> vec(mean(moments(θ),dims=1)) # 1Xg
    # weight
    Ωinv = θ -> inv(NeweyWest(moments(θ)))
    # objective
    obj = θ -> dot(m(θ),Ωinv(θ),m(θ))
    # do minimization
    θhat, objvalue, converged = fminunc(obj, θ)
    # derivative of average moments
    D = try
        ForwardDiff.jacobian(m, vec(θhat))'
    catch
        Calculus.jacobian(m, vec(θhat), :central)'
    end
    # moment contributions at estimate
    ms = moments(θhat)
    return θhat, objvalue, D, ms, Ωinv(θhat), converged
end

# called with no args runs this example
function gmm()
    x = [ones(10) 1:10]
    y = [1:2;3;4;5;6;7;8;9;10]
    β = zeros(2)
    # 3 moment conditions
    moments = β -> x.*(y - exp.(x*β))
    # first round consistent
    W = eye(2)
    βhat, objvalue, D, ms, converged = gmm(moments, β, W)
    βhat
end

max_stars_repo_name: Hiroakiyusheng/Econometrics

max_issues_repo_name: Hiroakiyusheng/Econometrics

max_forks_repo_name: Hiroakiyusheng/Econometrics


text:
export SpecialEuclidean

"""
$(TYPEDEF)

Rigid transform between two Pose2's, assuming (x,y,theta).

Related

Pose3Pose3, Point2Point2, MutablePose2Pose2Gaussian, DynPose2, InertialPose3
"""
struct Pose2Pose2{T <: IIF.SamplableBelief} <: IIF.AbstractRelativeRoots
  z::T
  # empty constructor
  Pose2Pose2{T}() where {T <: IIF.SamplableBelief} = new{T}()
  # regular constructor
  Pose2Pose2{T}(z1::T) where {T <: IIF.SamplableBelief} = new{T}(z1)
end
# convenience and default constructor
Pose2Pose2(z::T=MvNormal(zeros(3),LinearAlgebra.diagm([1.0;1.0;1.0]))) where {T <: IIF.SamplableBelief} = Pose2Pose2{T}(z)
Pose2Pose2(::UniformScaling) = Pose2Pose2(MvNormal(zeros(3),LinearAlgebra.diagm([1.0;1.0;1.0])))

getSample(cf::CalcFactor{<:Pose2Pose2}, N::Int=1) = (rand(cf.factor.z,N), )

function (cf::CalcFactor{<:Pose2Pose2})(meas,
                                        wxi,
                                        wxj  )
  #
  wXjhat = SE2(wxi)*SE2(meas)
  jXjhat = SE2(wxj) \ wXjhat
  return se2vee(jXjhat)
end

# NOTE, serialization support -- will be reduced to macro in future
# ------------------------------------

"""
$(TYPEDEF)
"""
mutable struct PackedPose2Pose2  <: IIF.PackedInferenceType
  datastr::String
  PackedPose2Pose2() = new()
  PackedPose2Pose2(x::String) = new(x)
end
function convert(::Type{Pose2Pose2}, d::PackedPose2Pose2)
  return Pose2Pose2(convert(SamplableBelief, d.datastr))
end
function convert(::Type{PackedPose2Pose2}, d::Pose2Pose2)
  return PackedPose2Pose2(convert(PackedSamplableBelief, d.z))
end

# FIXME, rather have separate compareDensity functions
function compare(a::Pose2Pose2,b::Pose2Pose2; tol::Float64=1e-10)
  return compareDensity(a.z, b.z)
  # TP = true
  # TP = TP && norm(a.z.μ-b.z.μ) < (tol + 1e-5)
  # TP = TP && norm(a.z.Σ.mat-b.z.Σ.mat) < tol
  # return TP
end

## Deprecated
# function Pose2Pose2(mean::Array{Float64,1}, cov::Array{Float64,2}, w::Vector{Float64})
#   @warn "Pose2Pose2(mu,cov,w) is deprecated in favor of Pose2Pose2(T(...)) -- use for example Pose2Pose2(MvNormal(mu, cov))"
#   Pose2Pose2(MvNormal(mean, cov))
# end


text:
"""
    cudnnSoftmaxForward(x; algo, mode, alpha)
    cudnnSoftmaxForward!(y, x; algo, mode, alpha, beta)

Return the softmax or logsoftmax of the input `x` depending on the `algo` keyword argument.
The `y` argument holds the result and it should be similar to `x` if specified. Keyword
arguments:

* `algo = (CUDA.math_mode()===CUDA.FAST_MATH ? CUDNN_SOFTMAX_FAST : CUDNN_SOFTMAX_ACCURATE)`: Options are `CUDNN_SOFTMAX_ACCURATE` which subtracts max from every point to avoid overflow, `CUDNN_SOFTMAX_FAST` which doesn't and `CUDNN_SOFTMAX_LOG` which returns logsoftmax.
* `mode = CUDNN_SOFTMAX_MODE_INSTANCE`: Compute softmax per image (N) across the dimensions C,H,W. `CUDNN_SOFTMAX_MODE_CHANNEL` computes softmax per spatial location (H,W) per image (N) across the dimension C.
* `alpha=1, beta=0` can be used for scaling, i.e. `y .= alpha * op(x1) .+ beta * y`
"""
cudnnSoftmaxForward, cudnnSoftmaxForward!

# Public methods
cudnnSoftmaxForward(x; o...) = cudnnSoftmaxForwardWithDefaults(x; o...)
cudnnSoftmaxForward!(y, x; o...) = cudnnSoftmaxForwardWithDefaults(x; y, o...)

# Private method
function cudnnSoftmaxForwardWithDefaults(
    x;
    y = similar(x),
    algo::cudnnSoftmaxAlgorithm_t = (CUDA.math_mode()===CUDA.FAST_MATH ? CUDNN_SOFTMAX_FAST : CUDNN_SOFTMAX_ACCURATE),
    mode::cudnnSoftmaxMode_t = CUDNN_SOFTMAX_MODE_INSTANCE,
    alpha::Real = 1,
    beta::Real = 0,
    format::cudnnTensorFormat_t = CUDNN_TENSOR_NCHW,
    xDesc::cudnnTensorDescriptor = cudnnTensorDescriptor(x; format),
    yDesc::cudnnTensorDescriptor = xDesc,
)
    @assert size(y) == size(x)
    T = eltype(x)
    alpha, beta = scalingParameter(T,alpha), scalingParameter(T,beta)
    cudnnSoftmaxForwardAD(x; algo, mode, alpha, xDesc, beta, yDesc, y)
end

# AD method
function cudnnSoftmaxForwardAD(x; algo, mode, alpha, xDesc, beta, yDesc, y)
    cudnnSoftmaxForward(handle(), algo, mode, alpha, xDesc, x, beta, yDesc, y)
    return y
end

# Deprecated methods
function cudnnSoftmaxForward(x::DenseCuArray{T,4}, y::DenseCuArray{T,4}; o...) where T
    @warn "`cudnnSoftmaxForward(x,y)` is deprecated, please use one of the methods in `@doc cudnnSoftmaxForward`." maxlog=1
    cudnnSoftmaxForward!(y, x; o...)
end


text:
abstract type JuliaConTalkType end
struct LightningTalk <: JuliaConTalkType end
struct SponsorTalk <: JuliaConTalkType end
struct Talk <: JuliaConTalkType end
struct Keynote <: JuliaConTalkType end
struct BoF <: JuliaConTalkType end
struct Minisymposium <: JuliaConTalkType end
struct Workshop <: JuliaConTalkType end
struct Experience <: JuliaConTalkType end
struct VirtualPoster <: JuliaConTalkType end
struct SocialHour <: JuliaConTalkType end
struct UnkownTalkType <: JuliaConTalkType end

const CONFERENCE_SCHEDULE_URL = "https://live.juliacon.org/agenda"

function json2df(conf)
    df = DataFrame(
        start = ZonedDateTime[],
        duration = CompoundPeriod[],
        title = String[],
        speaker = Vector{String}[],
        type = JuliaConTalkType[],
        url = String[],
        track = String[],
    )

    for day in conf["days"]
        date = Date(day["date"])
        for (track, talks) in day["rooms"]
            for talk in talks
                # parse duration
                tmp = split(talk["duration"], ':')
                dur = Hour(tmp[1]) + Minute(tmp[2])

                d = Dict(
                    :start => ZonedDateTime(DateTime(date, Time(talk["start"])), JULIACON_TIMEZONE),
                    :duration => dur,
                    :title => talk["title"],
                    :speaker => String[p["public_name"] for p in talk["persons"]],
                    :type => talktype_from_str(talk["type"]),
                    :url => talk["url"],
                    :track => track,
                )
                push!(df, d)
            end
        end
    end

    return df
end

function talktype_from_str(str)
    if str == "Talk"
        return Talk()
    elseif str == "Lightning talk"
        return LightningTalk()
    elseif str == "Sponsor Talk"
        return SponsorTalk()
    elseif str == "Keynote"
        return Keynote()
    elseif str == "Birds of Feather" || str == "BoF (45 mins)"
        return BoF()
    elseif str == "Minisymposium"
        return Minisymposium()
    elseif contains(str, "Workshop")
        return Workshop()
    elseif contains(str, "Experience")
        return Experience()
    elseif contains(str, "Virtual Poster")
        return VirtualPoster()
    elseif contains(str, "Social hour")
        return SocialHour()
    else
        return UnkownTalkType()
    end
end

string(x::LightningTalk) = "Lightning Talk"
string(x::SponsorTalk) = "Sponsor Talk"
string(x::Talk) = "Talk"
string(x::Keynote) = "Keynote"
string(x::BoF) = "Birds of Feather"
string(x::Minisymposium) = "Minisymposium"
string(x::Workshop) = "Workshop"
string(x::Experience) = "Experience"
string(x::VirtualPoster) = "Virtual Poster"
string(x::SocialHour) = "Social hour"
string(x::UnkownTalkType) = "Unknown"

abbrev(::Type{LightningTalk}) = "L"
abbrev(::Type{SponsorTalk}) = "S"
abbrev(::Type{Talk}) = "T"
abbrev(::Type{Keynote}) = "K"
abbrev(::Type{BoF}) = "BoF"
abbrev(::Type{Minisymposium}) = "M"
abbrev(::Type{Workshop}) = "W"
abbrev(::Type{Experience}) = "E"
abbrev(::Type{VirtualPoster}) = "P"
abbrev(::Type{SocialHour}) = "SH"
abbrev(::Type{UnkownTalkType}) = "U"

abbrev(x::JuliaConTalkType) = abbrev(typeof(x))

function _track2color(track::String)
    if track == "Red"
        return :red
    elseif track == "Green"
        return :green
    elseif track == "Purple"
        return :magenta
    elseif track == "Blue"
        return :blue
    elseif track == "JuMP Track"
        return 208 # orange, see https://github.com/KristofferC/Crayons.jl/blob/master/README.md
    else
        return :default
    end
end

is_schedule_json_available() = isfile(joinpath(CACHE_DIR, "schedule.json"))

"""
    get_all_tracks()

Returns a list of strings containing all track names.
"""
function get_tracks()
    cs = get_conference_schedule()
    return unique(cs.track)
end

"""
    print_legend(highlighting)

`highlighting` is a `Bool`.
"""
function print_legend(highlighting)
    println()
    if highlighting
        printstyled("Currently running talks are highlighted in ")
        printstyled("yellow"; color=:yellow)
        printstyled(".")
        println()
        println()
    end
    print(abbrev(Talk), " = Talk, ")
    print(abbrev(LightningTalk), " = Lightning Talk, ")
    print(abbrev(SponsorTalk), " = Sponsor Talk, ")
    println(abbrev(Keynote), " = Keynote, ")
    print(abbrev(Workshop), " = Workshop, ")
    print(abbrev(Minisymposium), " = Minisymposium, ")
    println(abbrev(BoF), " = Birds of Feather, ")
    print(abbrev(Experience), " = Experience, ")
    println(abbrev(VirtualPoster), " = Virtual Poster")
    println()
    println("Check out $(CONFERENCE_SCHEDULE_URL) for more information.")
end

"""
    get_conference_schedule(; speaker=nothing)

Get the conference schedule as a DataFrame.
On first call, the schedule is downloaded from Pretalx and cached for further usage.
`speaker` can be a string identifying the speaker to filter the schedule.
"""
function get_conference_schedule(; speaker=nothing)
    isassigned(jcon) || update_schedule()
    # filter for speaker
    jcon_filt = filter(jcon[]) do talk
        return isnothing(speaker) || any(contains.(talk.speaker, speaker))
    end
    return jcon_filt
end

"""
    update_schedule(; verbose=false, ignore_timeout=false)

Explicitly trigger a schedule update according to the specified CACHE_MODE.
"""
function update_schedule(; verbose=false, notimeout=false)
    local file
    to = TimerOutput()
    verbose && @info "Cache mode: $CACHE_MODE"
    usecache = CACHE_MODE != :NEVER
    download_dir = usecache ? CACHE_DIR : mktempdir()

    if CACHE_MODE != :ALWAYS
        verbose && @info "Downloading $(default_json_url()) to $download_dir"
        if usecache && !isdir(CACHE_DIR)
            verbose && @info "Cache directory $CACHE_DIR created."
            mkpath(CACHE_DIR)
        end

        timeout =
            (usecache && is_schedule_json_available()) ? (!notimeout ? TIMEOUT : Inf) : Inf
        verbose && @info "Timeout set to $timeout seconds."

        try
            @timeit to "download" file = download(
                default_json_url(), joinpath(download_dir, "schedule.json.tmp"); timeout
            )
            file = mv(
                joinpath(download_dir, "schedule.json.tmp"),
                joinpath(download_dir, "schedule.json");
                force=true,
            )
        catch err
            if usecache
                @warn "Download failed or timed out. Falling back to cached schedule (might be stale). " *
                      "You can try forcing matters with JuliaCon.update_schedule(notimeout=true) or " *
                      "skipping the update altogether via JuliaCon.set_cachemode(:ALWAYS)."
                file = joinpath(CACHE_DIR, "schedule.json")
            else
                error(
                    "Download failed. Not using the cache due to CACHE_MODE = $CACHE_MODE."
                )
            end
        end
    else
        verbose && @info "Loading cached schedule.json"
        is_schedule_json_available() || error(
            "Can't find cached schedule.json. Not downloading due to CACHE_MODE = $CACHE_MODE.",
        )
        file = joinpath(CACHE_DIR, "schedule.json")
    end

    @timeit to "parse2json" data = JSON.parsefile(file)
    @timeit to "json2df" jcon[] = json2df(data["schedule"]["conference"])

    verbose && @info string("Timings:\n", to)
    return nothing
end

function get_running_talks(; now=default_now())
    jcon = get_conference_schedule()

    running_talks = filter(jcon; view=true) do talk
        start_time = talk.start
        end_time = start_time + talk.duration
        return start_time <= astimezone(now, JULIACON_TIMEZONE) < end_time
    end
    return running_talks
end

function _print_running_talks(running_talks; now=default_now())
    nrow(running_talks) > 0 || return nothing
    # println()
    # println(Dates.format(default_now(), "HH:MM dd-mm-YYYY"))
    for talk in eachrow(running_talks)
        println()
        printstyled(_add_track_emoji(string(talk.track)); bold=true, color=_track2color(talk.track))
        println()
        println("\t", talk.title, " (", string(talk.type), ")")
        println("\t", "├─ ", _speakers2str(talk.speaker))
        println("\t", "└─ ", talk.url)
    end
    println("\n")
    println("(Full schedule: $(CONFERENCE_SCHEDULE_URL))")
    return nothing
end

function now(::Val{:text}; now)
    running_talks = get_running_talks(; now=now)
    str = ""
    if !isnothing(running_talks)
        for talk in eachrow(running_talks)
            str *= """
            $(_add_track_emoji(string(talk.track)))
            \t$(talk.title) ($(string(talk.type)))
            \t├─ $(JuliaCon._speakers2str(talk.speaker))
            \t└─ $(talk.url)
            """
        end
    end
    str *= "\n(Full schedule: $(CONFERENCE_SCHEDULE_URL))"
    return str
end

function now(::Val{:terminal}; now)
    running_talks = get_running_talks(; now=now)
    _print_running_talks(running_talks; now=now)
    return nothing
end

# A dispatcher for the `now` methods. Default to terminal output.
now(; now=default_now(), output=:terminal) = JuliaCon.now(Val(output); now=now)

function _speakers2str(speaker::Vector{String})
    if length(speaker) <= 3
        return join(speaker, ", ")
    else
        return string(join(speaker[1:3], ", "), " et al.")
    end
end

function _get_today_tables(;
    now=default_now(), speaker=nothing, track=nothing, terminal_links=TERMINAL_LINKS,
    highlighting=true, text_highlighting=false,
    )
    jcon = get_conference_schedule(speaker=speaker)

    today_start_utc = ZonedDateTime(DateTime(Date(now), Time("00:00")), JULIACON_TIMEZONE)
    today_end_utc = ZonedDateTime(DateTime(Date(now) + Day(1), Time("00:00")), JULIACON_TIMEZONE)

    talks_today = filter(jcon; view=true) do talk
        # talk is in the requested track (default: any)
        if !isnothing(track) && talk.track != track
            return false
        end

        # talk starts "today" (local time)
        return today_start_utc <= talk.start < today_end_utc
    end

    # no talks today -> exit
    nrow(talks_today) > 0 || return (nothing, nothing, nothing)

    # create talk tables
    tracks = String[]
    tables = Matrix{Union{String,URLTextCell}}[]
    highlighters = Union{Nothing,Highlighter}[]

    # for each track
    for track_grp in groupby(talks_today, :track)
        # build talk-data matrix
        data = Matrix{Union{String,URLTextCell}}(undef, nrow(track_grp), 4)
        for (i, talk) in enumerate(eachrow(track_grp))
            data[i, 1] = Dates.format(astimezone(talk.start, timezone(now)), "HH:MM")
            data[i, 2] = terminal_links ? URLTextCell(talk.title, talk.url) : talk.title
            data[i, 3] = JuliaCon.abbrev(talk.type)
            data[i, 4] = JuliaCon._speakers2str(talk.speaker)
        end

        h_running = Highlighter((data, m, n) -> false, crayon"yellow")
        if highlighting
            for (i, talk) in enumerate(eachrow(track_grp))
                start_time = talk.start
                end_time = talk.start + talk.duration
                if start_time <= astimezone(now, JULIACON_TIMEZONE) < end_time
                    if text_highlighting
                        if data[i, 2] isa URLTextCell
                            data[i, 2].text = string("> ", data[i, 2].text)
                        else
                            data[i, 2] = string("> ", data[i, 2])
                        end
                    end

                    h_running = Highlighter((data, m, n) -> m == i, crayon"yellow")
                end
            end
        end

        push!(tracks, first(track_grp).track)
        push!(tables, data)
        push!(highlighters, h_running)
    end

    @assert length(tables) == length(highlighters) == length(tracks)
    return (tracks, tables, highlighters)
end

# A dispatcher for the `today` methods. Defaults to terminal output.
"""     today(; speaker)

Prints the schedule of today.
`speaker` can be a string identifying the speaker to filter the schedule.
`track` can be string used to filter for a track. See `JuliaCon.get_tracks()` for possible options.
"""
function today(;
    now=default_now(),
    speaker=nothing,
    track=nothing,
    terminal_links=TERMINAL_LINKS,
    output=:terminal, # can take the :text value to output a Vector{String}
    highlighting=true,
    legend=true,
)
    return today(Val(output); now, speaker, track, terminal_links, highlighting, legend)
end

function today(::Val{:terminal}; now, speaker, track, terminal_links, highlighting=true, legend=true)
    tracks, tables, highlighters = _get_today_tables(; now, speaker, track, terminal_links, highlighting)
    isnothing(tables) && return nothing

    pretty_print_results(now, tracks, legend, highlighting, tables, highlighters)
    return nothing
end

function _add_track_emoji(track::AbstractString)
    if track == "Red"
        return "🍎 Red"
    elseif track == "Green"
        return "🍏 Green"
    elseif track == "Blue"
        return "🔷 Blue"
    elseif track == "Purple"
        return "💜 Purple"
    elseif track == "BoF/Mini Track"
        return "🕊  BoF/Mini Track"
    elseif track == "JuMP Track"
        return "🔸 JuMP Track"
    else
        return track
    end
end

function today(::Val{:text}; now, speaker, track, terminal_links, highlighting=true, legend=true)
    tracks, tables, highlighters = _get_today_tables(; now, speaker, track, terminal_links, highlighting, text_highlighting=highlighting)
    isnothing(tables) && return nothing

    results_to_string(now, tracks, legend, highlighting, tables, highlighters)
end

"""     tomorrow(; speaker)

Prints the schedule of tomorrow.
`speaker` can be a string identifying the speaker to filter the schedule.
`track` can be string used to filter for a track. See `JuliaCon.get_tracks()` for possible options.
"""
function tomorrow(;
    now=default_now(),
    speaker=nothing,
    track=nothing,
    terminal_links=TERMINAL_LINKS,
    output=:terminal, # can take the :text value to output a Vector{String},
    legend=true
)
    return today(Val(output); now = now + Dates.Day(1), speaker, track, terminal_links, highlighting = false,
                 legend)
end

"""
    talks_by(speaker; output=:terminal, legend=false)

Prints all talks of a speaker identified by the (sub-)string `speaker`
"""
function talks_by(speaker; output=:terminal, legend=false)
    df = get_conference_schedule()
    # strip of time and filter for pure days
    days = unique(map(x -> Dates.yearmonthday(x), df.start))

    # list of ZoneDateTime for all JuliaCon days
    all_juliacon_dates = map(d -> ZonedDateTime(d..., JuliaCon.LOCAL_TIMEZONE), days)
    t(d, legend) = _get_today_tables(; now = d, speaker, legend)

    _print_talks_by(all_juliacon_dates, speaker, legend, Val(output))
end

function _print_talks_by(all_juliacon_dates, speaker, legend, ::Val{:text})
    str = []
    for d in all_juliacon_dates
        tracks, tables, highlighters = _get_today_tables(; now=d, speaker)
        isnothing(tables) && continue
        s = results_to_string(d, tracks, legend, true, tables, highlighters)
        append!(str, s)
    end
    return str
end

function _print_talks_by(all_juliacon_dates, speaker, legend, ::Val{:terminal})
    for d in all_juliacon_dates
        tracks, tables, highlighters = _get_today_tables(; now=d, speaker)
        isnothing(tables) && continue
        pretty_print_results(d, tracks, legend, true, tables, highlighters)
    end
    return nothing
end


text:
using DataFrames, Plots
pyplot()

ProjDir = dirname(@__FILE__)
cd(ProjDir) do

  function euler_ex71(f::Function, y0::Float64, x::FloatRange{Float64})
    y = Vector{Real}(length(x))
    y[1] = y0
    h = x.step/x.divisor
    for i in 2:length(x)
      k0 = h * f(x[i-1], y[i-1])
      k1 = h * f(x[i-1] + h, y[i-1] + k0)
      y[i] = y[i-1] + 1/2*(k0 + k1)
    end
    columns=[]
    push!(columns, collect(x))
    push!(columns, y)
    df = DataFrame(columns, [:x, :y])
  end

  f(x,y) = (x+y)/x
  fy(x) = x*(1+log(x/2))

  x1 = 2.0:0.25:2.5
  x2 = 2.0:0.1:2.5
  x3 = 1.8:0.1:2.6
  df1 = euler_ex71(f, 2.0, x1)
  df2 = euler_ex71(f, 2.0, x2)

  df1 |> display
  println()
  df2 |> display
  println()

  global p2

  p1 = plot(x3, fy.(x3), ylims=(1.8, 3.5))
  plot!(p1, df1[:x], df1[:y], line=(:dash, 1), marker=(:plus), color=:red)
  plot!(p1, df2[:x], df2[:y], line=(:dash, 1), marker=(:circle), color=:red)
  savefig("euler_ex7.2.png")
  plot(p1, layout=(1,1))
  gui()
end


text:
mutable struct MarkovChain{T <: EpidemicModel}
  iterations::Int64
  events::Vector{Events{T}}
  transmission_network::Vector{TransmissionNetwork}
  risk_parameters::Vector{RiskParameters{T}}
  log_posterior::Vector{Float64}
  cov::OnlineStats.CovMatrix

  function MarkovChain(events::Events{T},
                       tn::TransmissionNetwork,
                       rp::RiskParameters{T},
                       lp::Float64) where  T <: EpidemicModel
    return new{T}(0, [events], [tn], [rp], [lp], OnlineStats.CovMatrix())
  end

  function MarkovChain{T}() where  T <: EpidemicModel
    return new{T}(0, Events{T}[], TransmissionNetwork[], RiskParameters{T}[], Float64[], OnlineStats.CovMatrix())
  end
end

function Base.length(x::MarkovChain{T}) where T <: EpidemicModel
  return x.iterations
end

function Base.show(io::IO, x::MarkovChain{T}) where T <: EpidemicModel
  return print(io, "$T model Markov chain (iterations = $(length(x)))")
end

function TransmissionNetworkDistribution(x::MarkovChain; burnin::Int64=0, thin::Int64=1)
  return TNDistribution(x.transmission_network[burnin:thin:end])
end


text:
module JSON2Ex
using JSON2

function JSON2.write(io::IO, obj::AbstractString; kwargs...)
    if startswith(obj, "{\"type\":") || startswith(obj, "{\"coordinates\":")
        Base.write(io, obj)
    else
        Base.write(io, '"')
        if JSON2.needescape(obj)
            bytes = codeunits(obj)
            for i = 1:length(bytes)
                @inbounds b = JSON2.ESCAPECHARS[bytes[i] + 0x01]
                Base.write(io, b)
            end
        else
            Base.write(io, obj)
        end
        Base.write(io, '"')
    end
    return
end

end

max_stars_repo_name: huaxk/rlservice

max_issues_repo_name: huaxk/rlservice

max_forks_repo_name: huaxk/rlservice


text:
"""運動学

"""
module Kinematics

export Phi_0
export Phi_1
export Phi_2
export Phi_3
export Phi_4

export Arm
export Ξ

function Phi_0!(q::Vector{Float64}, ξ::Float64, out::Vector{Float64})
    ccall(
        (:Phi_0, "o/soft_robot/derivation_of_kinematics/derived/so/Phi_0.so"),
        Cvoid,
        (Ptr{Cdouble}, Cdouble, Ptr{Cdouble}),
        q, ξ, out
    )
end

"""グローバル位置"""
function Phi_0(q::Vector{Float64}, ξ::Float64)
    Z = Vector{Float64}(undef, 3)
    Phi_0!(q, ξ, Z)
    return Z
end

function Phi_1!(q::Vector{Float64}, ξ::Float64, out::Vector{Float64})
    ccall(
        (:Phi_1, "o/soft_robot/derivation_of_kinematics/derived/so/Phi_1.so"),
        Cvoid,
        (Ptr{Cdouble}, Cdouble, Ptr{Cdouble}),
        q, ξ, out
    )
end

"""グローバル位置"""
function Phi_1(q::Vector{Float64}, ξ::Float64)
    Z = Vector{Float64}(undef, 3)
    Phi_1!(q, ξ, Z)
    return Z
end

function Phi_2!(q::Vector{Float64}, ξ::Float64, out::Vector{Float64})
    ccall(
        (:Phi_2, "o/soft_robot/derivation_of_kinematics/derived/so/Phi_2.so"),
        Cvoid,
        (Ptr{Cdouble}, Cdouble, Ptr{Cdouble}),
        q, ξ, out
    )
end

"""グローバル位置"""
function Phi_2(q::Vector{Float64}, ξ::Float64)
    Z = Vector{Float64}(undef, 3)
    Phi_2!(q, ξ, Z)
    return Z
end

function Phi_3!(q::Vector{Float64}, ξ::Float64, out::Vector{Float64})
    ccall(
        (:Phi_3, "o/soft_robot/derivation_of_kinematics/derived/so/Phi_3.so"),
        Cvoid,
        (Ptr{Cdouble}, Cdouble, Ptr{Cdouble}),
        q, ξ, out
    )
end

"""グローバル位置"""
function Phi_3(q::Vector{Float64}, ξ::Float64)
    Z = Vector{Float64}(undef, 3)
    Phi_3!(q, ξ, Z)
    return Z
end

function Phi_4!(q::Vector{Float64}, ξ::Float64, out::Vector{Float64})
    ccall(
        (:Phi_4, "o/soft_robot/derivation_of_kinematics/derived/so/Phi_4.so"),
        Cvoid,
        (Ptr{Cdouble}, Cdouble, Ptr{Cdouble}),
        q, ξ, out
    )
end

"""グローバル位置"""
function Phi_4(q::Vector{Float64}, ξ::Float64)
    Z = Vector{Float64}(undef, 3)
    Phi_4!(q, ξ, Z)
    return Z
end

const Ξ = Vector(0:0.01:1)

"""アームのディスク位置を計算"""
function Arm(q::Vector{T}, section_n::Int64) where T
    n = length(Ξ)
    Phis = Vector{Vector{T}}(undef, section_n*n)
    for (i, ξ) in enumerate(Ξ)
        Phis[i] = Phi_0(q, ξ)
        if section_n > 1
            Phis[i+n] = Phi_1(q, ξ)
        elseif section_n > 2
            Phis[i+2n] = Phi_2(q, ξ)
        elseif  section_n > 3
            Phis[i+3n] = Phi_3(q, ξ)
        elseif  section_n > 4
            Phis[i+4n] = Phi_4(q, ξ)
        end
        #println(Phis[i], Phis[i+n], Phis[i+2n])
    end
    return Phis
end

# """アームのディスク位置を計算"""
# function Arm(q::Vector{T}, m::Int64) where T
#     n = length(Ξ)
#     Phis = Vector{Vector{T}}(undef, n)
#     for (i, ξ) in enumerate(Ξ)
#         if m == 0
#             Phis[i] = Phi_0(q, ξ)
#         elseif m == 1
#             Phis[i] = Phi_1(q, ξ)
#         else
#             Phis[i] = Phi_2(q, ξ)
#         end

#     end
#     return Phis
# end

end

# using .Kinematics

# q = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
# p0 = Phi_0(q, 0.5)
# p1 = Phi_1(q, 0.5)
# p2 = Phi_2(q, 0.5)

# #Arm(q, 0)
# println(p0)
# println(p1)
# println(p2)


text:
include(pwd() * "/examples/boiler_plate.jl")
# include(pwd() * "/src/bigfileofstuff.jl")
using QuickVizExample

jldfile = jldopen("example.jld2")
##
# reconstruct grid
Ω = jldfile["Ω"]
elements = jldfile["elements"]
polynomialorder = jldfile["polynomialorder"]
periodicity = (true, true, false)
dgΩ = DiscontinuousSpectralElementGrid(Ω, periodicity = periodicity,
    elements = elements, polynomialorder = polynomialorder,
    mpicomm=MPI.COMM_WORLD, array = Array)
# construct gridhelpers
gridhelper = GridHelper(dgΩ)
x, y, z = coordinates(dgΩ)
ϕ = ScalarField(copy(x), gridhelper)

# construct new grid
xnew = range(Ω[1][1], Ω[1][2], length = 128)
ynew = range(Ω[2][1], Ω[2][2], length = 128)
znew = range(Ω[3][1], Ω[3][1], length = 1)

# load data
u = jldfile["velocity_u"]
v = jldfile["velocity_v"]
η = jldfile["scalar_η"]
c = jldfile["scalar_c"]

# interpolate each moment in time to a new grid
nt = length(u)
ut = zeros(length(xnew), length(ynew), nt)
vt = zeros(length(xnew), length(ynew), nt)
ηt = zeros(length(xnew), length(ynew), nt)
ct = zeros(length(xnew), length(ynew), nt)
tic = time()
trange = collect(1:nt)
for i in 1:nt
    ϕ .= u[i]
    ut[:, :, i] .= view(ϕ(xnew, ynew, znew), :, :, 1)
    ϕ .= v[i]
    vt[:, :, i] .= view(ϕ(xnew, ynew, znew), :, :, 1)
    ϕ .= η[i]
    ηt[:, :, i] .= view(ϕ(xnew, ynew, znew), :, :, 1)
    ϕ .= c[i]
    ct[:, :, i] .= view(ϕ(xnew, ynew, znew), :, :, 1)
end
toc = time()
println("interpolation time is $(toc - tic) seconds")
# visualize (z-axis is time here)
states = [ut, vt, ηt, ct]
statenames = ["u", "v", "η", "c"]
scene = volumeslice(states, statenames = statenames)
visualize(states)
visualize([ut[:,:,1]])

##
seconds = 5
fps = 30
frames = round(Int, fps * seconds )
record(scene, pwd() * "/example.mp4"; framerate = fps) do io
    for i = 1:frames
        sleep(1/fps)
        recordframe!(io)
    end
end


text:
import Base

# Range

const range = Base.range

# Repeat

const repeat = Base.repeat

# First

const first = Base.first

# Last

const last = Base.last

# LastIndex

const lastindex = Base.lastindex

# Empty

const empty = Base.isempty

# Reverse

const reverse = Base.Iterators.reverse

# Zip

zip(c::Tuple) = zip(c...)

zip(c...) = Base.zip(c...)

# Unzip

unzip(c) = Base.map(f -> getfield.(c, f), fieldnames(eltype(c)))

# Enumerate

enumerate(c) = Base.enumerate(c)

# Map

map(f::Function) = c -> Base.map(f, c)

mapi(f::Function) = c -> begin
    n = length(c)
    A = Vector{Any}(undef, n)
    for i = 1:n
        A[i] = f(i, c[i])
    end
    (A...,)
end

# Filter

filter(f::Function) = c -> Base.filter(f, c)

# All

all(f::Function) = c -> c |> map(f) |> Base.all

# Any

any(f::Function) = c -> c |> map(f) |> Base.any

# Take

take(n::Int) = c -> Base.Iterators.take(c, n)

# Take While

takewhile(f::Function) = c -> Base.Iterators.takewhile(f, c)

# Skip

skip(n::Int) = c -> Base.Iterators.drop(c, n)

# Skip While

skipwhile(f::Function) = c -> Base.Iterators.dropwhile(f, c)

# Reduce

reduce(f::Function) = c -> Base.reduce(f, c)

# Fold Right

foldright(f::Function) = c -> Base.foldr(f, c)

# Fold Left

foldleft(f::Function) = c -> Base.foldl(f, c)

# Sum

sum(c) = c |> reduce(+)

# Mul

mul(c) = c |> reduce(*)

# Max

max(c) = Base.Iterators.maximum(c)

max(f::Function) = c -> Base.Iterators.maximum(f, c)

# Min

min(c) = Base.Iterators.minimum(c)

min(f::Function) = c -> Base.Iterators.minimum(f, c)

# Avg

avg(c) = sum(c) / length(c)

avg(f::Function) = c -> c |> map(f) |> avg

# Find

find(f) = x -> x |> filter(f) |> first

# Join

join(delim) = c -> Base.join(c, delim)

# Concat

concat(c::Tuple) = concat(c...)

concat(c...) = Base.vcat(c...)

# Group By

groupby(f::Function) = c -> begin
    out = Dict()
    for e ∈ c
        k = f(e)
        l = get(out, k, nothing)
        if l === nothing
            l = []
            out[k] = l
        end
        push!(l, e)
    end
    return out |> collect |> map(e -> e[2])
end

# Merge

merge(f::Function) = c -> begin
end

# Apply

apply(f::Function) = c -> begin
    for e ∈ c
        f(e)
    end
end

# Distinct

distinct(f::Function) = c -> begin
    out = Dict()
    for e ∈ c
        k = f(e)
        l = get(out, k, nothing)
        if l === nothing
            out[k] = e
        end
    end
    return out |> collect |> map(e -> e[2])
end


text:
# Equality

# equals() now depends on == instead
# of the other way round.
function ==(x::Decimal, y::Decimal)
    a = normalize(x)
    b = normalize(y)
    a.c == b.c && a.q == b.q && a.s == b.s
end

Base.iszero(x::Decimal) = iszero(x.c)

function <(x::Decimal, y::Decimal)
    # return early on zero
    if iszero(x) && iszero(y)
        return false
    end

    signdiff = y.s - x.s
    if signdiff == 1
        return false
    elseif signdiff == -1
        return true
    end

    diff = y - x

    farther_from_0 = diff.c > 0 || (iszero(diff.c) && diff.q > 0)

    if diff.s == 1
        return !farther_from_0
    else
        return farther_from_0
    end
end


text:
loss(ŷ, y) = logitcrossentropy(ŷ, y)

num_params(model) = sum(length, Flux.params(model))

round4(x) = round(x, digits = 4)

function accuracy(preds, labels)
    acc = sum(onecold(preds |> cpu) .== onecold(labels |> cpu))
    return acc
end

function get_data(args)
    xtrain, ytrain = MLDatasets.MNIST.traindata(Float32)
    xtest, ytest = MLDatasets.MNIST.testdata(Float32)

    xtrain = reshape(xtrain, 28, 28, 1, :)
    xtest = reshape(xtest, 28, 28, 1, :)

    ytrain, ytest = onehotbatch(ytrain, 0:9), onehotbatch(ytest, 0:9)

    train_loader = DataLoader(
        (xtrain, ytrain),
        batchsize = args.batchsize,
        shuffle = true,
        partial = false,
    )
    test_loader = DataLoader((xtest, ytest), batchsize = args.batchsize, partial = false)

    # Fashion mnist for uncertainty testing
    xtrain, ytrain = MLDatasets.FashionMNIST.traindata(Float32)
    xtest, ytest = MLDatasets.FashionMNIST.testdata(Float32)

    xtrain = reshape(xtrain, 28, 28, 1, :)
    xtest = reshape(xtest, 28, 28, 1, :)

    ytrain, ytest = onehotbatch(ytrain, 0:9), onehotbatch(ytest, 0:9)

    ood_train_loader = DataLoader(
        (xtrain, ytrain),
        batchsize = args.batchsize,
        shuffle = true,
        partial = false,
    )
    ood_test_loader =
        DataLoader((xtest, ytest), batchsize = args.batchsize, partial = false)

    return train_loader, test_loader, ood_train_loader, ood_test_loader
end

function ensembles_evaluation(args, loader, model, device)
    l = [0.0f0 for x = 1:args.ensemble_size]
    acc = [0 for x = 1:args.ensemble_size]
    ece_list = [0.0f0 for x = 1:args.ensemble_size]
    ntot = 0
    mean_l = 0
    mean_acc = 0
    mean_ece = 0
    for (x, y) in loader
        x = repeat(x, 1, 1, 1, args.ensemble_size)
        x, y = x |> device, y |> device
        # Perform the forward pass
        ŷ = model(x)
        # Reshape the predictions into [classes, batch_size, ensemble_size
        reshaped_ŷ = reshape(ŷ, size(ŷ)[1], args.batchsize, args.ensemble_size)
        # Loop through each model's predictions
        for ensemble = 1:args.ensemble_size
            model_predictions = reshaped_ŷ[:, :, ensemble]
            # Calculate individual loss
            l[ensemble] += loss(model_predictions, y) * size(model_predictions)[end]
            acc[ensemble] += accuracy(model_predictions, y)
            ece_list[ensemble] +=
                expected_calibration_error(model_predictions, onecold(y)) * args.batchsize
        end
        # Get the mean predictions
        mean_predictions = mean(reshaped_ŷ, dims = ndims(reshaped_ŷ))
        mean_predictions = dropdims(mean_predictions, dims = ndims(mean_predictions))
        mean_l += loss(mean_predictions, y) * size(mean_predictions)[end]
        mean_acc += accuracy(mean_predictions, y)
        mean_ece +=
            expected_calibration_error(mean_predictions, onecold(y)) * args.batchsize
        ntot += size(mean_predictions)[end]
    end
    # Normalize the loss
    losses = round4.(l ./ ntot)
    acc = round4.((acc ./ ntot) .* 100)
    ece_list = round4.(ece_list ./ ntot)
    # Calculate mean loss
    mean_l = round4(mean_l / ntot)
    mean_acc = round4((mean_acc / ntot) .* 100)
    mean_ece = round4(mean_ece / ntot)

    # Print the per ensemble mode loss and accuracy
    for ensemble = 1:args.ensemble_size
        @info (format(
            "Model {} Loss: {} Accuracy: {} ECE: {}",
            ensemble,
            losses[ensemble],
            acc[ensemble],
            ece_list[ensemble],
        ))
    end
    @info (format(
        "Mean Loss: {} Mean Accuracy: {} Mean ECE: {}",
        mean_l,
        mean_acc,
        mean_ece,
    ))
    @info "========================="
    return nothing
end

function monte_carlo_evaluation(args, loader, model, device)
    ntot = 0
    mean_l = 0
    mean_acc = 0
    mean_ece = 0
    mean_entropy = 0
    for (x, y) in loader
        predictions = []
        x, y = x |> device, y |> device

        # Loop through each model's predictions
        for ensemble = 1:args.sample_size
            model_predictions = model(x)
            push!(predictions, model_predictions)
        end
        # Get the mean predictions
        predictions = Flux.batch(predictions)
        mean_predictions = mean(predictions, dims = ndims(predictions))
        mean_predictions = dropdims(mean_predictions, dims = ndims(mean_predictions))
        mean_l += loss(mean_predictions, y) * size(mean_predictions)[end]
        mean_acc += accuracy(mean_predictions, y)
        mean_ece +=
            expected_calibration_error(mean_predictions, onecold(y)) * args.batchsize
        mean_entropy += mean(calculate_entropy(mean_predictions)) * args.batchsize
        ntot += size(mean_predictions)[end]
    end
    # Calculate mean loss
    mean_l = mean_l / ntot |> round4
    mean_acc = mean_acc / ntot * 100 |> round4
    mean_ece = mean_ece / ntot |> round4
    mean_entropy = mean_entropy / ntot |> round4

    # Print the per ensemble mode loss and accuracy
    @info (format(
        "Mean Loss: {} Mean Accuracy: {} Mean ECE: {} Mean Entropy: {}",
        mean_l,
        mean_acc,
        mean_ece,
        mean_entropy,
    ))
    @info "========================="
    return nothing
end


text:
module Utils

include("Utils/mutref.jl")
include("Utils/thousand.jl")
include("Utils/unionfind.jl")
include("Utils/heap.jl")

end

max_stars_repo_name: EQt/graphidx

max_issues_repo_name: EQt/graphidx

max_forks_repo_name: EQt/graphidx