"""
`module DB`

## Overview

This module implements a "database" for storing a precomputed LSQ system. This
is useful for fitting e.g. NBodyIPs since it allows one to precompute the <basis,
data> inner products which are very expensive, and then quickly construct LSQ
systems from them using e.g., many variants of weighting and regularisation.

A db called (e.g.) 'lsqdb' consists of two files:
 * `lsqdb_info.json` : this stores a list of basis functions and a list
 of configurations with attached observations (normally DFT data)
 * `lsqdb_kron.h5` : this stores all the inner products <basis, data>
 as a single large matrix. The observations and basis functions in
 `lsqdb_info` store the corresponding column and row indices.

## More details and remarks

### INFO

* BASIS: The basis is (for the time being) simply represented as a list (Vector) of
JuLIP.AbstractCalculator. These are stored as `INFO["basis"]`.

* DATA: Each "datum" is an atomistic configuration, which must have a
"configtype", stored as a `Dat`.
Informally two pieces of data with the same configtype should represent
"similar" kinds of configurations. `data = INFO["data"]` is a Dictionary
where the keys are the configtypes.

### (De-)Vectorising Data

Forces in `JuLIP` are represented as `Vector{SVector{...}}`, which is the
same memory layout as a 3 x N matrix (`JuLIP.mat`), which can can then be
vectorised (`[:]`), and this vectoriation is readily undone again.
Analogously, any data must be stored in such a vectorised format. This is
achieved (e.g. for forces) via
* `vec_obs(::Val{:F}, F) -> Vector`
* `devec_obs(::Val{:F}, Fvec) -> Vector{SVector{...}}`
or equivalently
* `vec_obs("F", F) -> Vector`
* `devec_obs("F", Fvec) -> Vector{SVector{...}}`
(the `Val` versions are used for performance optimisation)
"""
module DB

using Base.Threads:          SpinLock, nthreads
using StaticArrays:          SVector
using JuLIP:                 AbstractCalculator, AbstractAtoms, Atoms,
                             save_json, load_json, decode_dict
using IPFitting:        Dat, LsqDB, basis, eval_obs, observations,
                             observation, vec_obs, devec_obs,
                             tfor_observations
using IPFitting.Data:   configtype
using HDF5:                  h5open, read

import Base: flush, append!, union

export LsqDB, info, configtypes

const KRONFILE = "_kron.h5"
const INFOFILE = "_info.json"

"""
`dbpath(db::LsqDB)` : return the absolute path to the database files, not
including the 'kron.h5' and 'info.json' endings.
"""
dbpath(db::LsqDB) = db.dbpath

infofile(dbpath::AbstractString) = dbpath * INFOFILE

kronfile(dbpath::AbstractString) = dbpath * KRONFILE

# ------------ Save and load the info file

"if the file exists, append a random string to avoid overwriting"
function _backupfile(fname)
   if isfile(fname)
      fnew = fname * "." * String(rand('a':'z', 5))
      @warn("The file $fname already exists. It will be renamed to $fnew to avoid overwriting.")
      run(`mv $fname $fnew`)
   end
   return nothing
end

function load_info(dbpath::String)
   dbinfo = load_json(infofile(dbpath))
   basis = decode_dict.(dbinfo["basis"])
   configs = Dat.(dbinfo["configs"])   # here we already know the type
   return basis, configs
end

function save_info(dbpath::String, db)
   _backupfile(infofile(dbpath))
   save_json(infofile(dbpath),
             Dict("basis" => Dict.(db.basis),
                  "configs" => Dict.(db.configs))
            )
   return nothing
end

save_info(db) = save_info(dbpath(db), db)

# ----------- Save and load the KRON file

"save a single matrix to HDF5"
_savemath5(A, fname) =
   h5open(fname, "w") do fid
      fid["A"] = A
      nothing
   end

"load a single matrix from HDF5"
_loadmath5(fname) =
   h5open(fname, "r") do fid
      read(fid["A"])
   end

function save_kron(dbpath, db)
   _backupfile(kronfile(dbpath))
   _savemath5(db.Ψ, kronfile(dbpath))
end

save_kron(db) = save_kron(dbpath(db), db)

load_kron(dbpath::String; mmap=false) = _loadmath5(kronfile(dbpath))

load_kron(db::LsqDB; mmap=false) = load_kron(dbpath(db); mmap=mmap)

function LsqDB(dbpath::AbstractString; mmap=false)
   basis, configs = load_info(dbpath)
   Ψ = load_kron(dbpath; mmap=mmap)
   return LsqDB(basis, configs, Ψ, dbpath)
end

"""
for the time being, this just checks that the db director and db name are
admissible. In the future, this should initialise the DB files.
"""
function initdb(basedir, dbname)
   @assert !('/' in dbname)
   @assert isdir(basedir)
   dbpath = joinpath(basedir, dbname)
   initdb(dbpath)
   return nothing
end

function initdb(dbpath)
   # TODO: seems this fails to detect an existing database?
   @assert !isfile(infofile(dbpath))
   @assert !isfile(kronfile(dbpath))
   # check that we can actually create and delete this file
   save_json(infofile(dbpath), Dict("basis" => [], "data" => []))
   rm(infofile(dbpath))
   _savemath5(rand(5,5), kronfile(dbpath))
   rm(kronfile(dbpath))
   return nothing
end

function flush(db::LsqDB)
   save_info(db)
   save_kron(db)
   return nothing
end

function LsqDB(dbpath::AbstractString,
               basis::AbstractVector{<: AbstractCalculator},
               configs::AbstractVector{Dat};
               verbose=true,
               maxnthreads=nthreads())
   # assign indices, count observations and allocate a matrix
   Ψ = _alloc_lsq_matrix(configs, basis)
   # create the struct where everything is stored
   db = LsqDB(basis, configs, Ψ, dbpath)
   # parallel assembly of the LSQ matrix
   tfor_observations( configs,
      (n, okey, cfg, lck) -> safe_append!(db, lck, cfg, okey),
      msg = "Assemble LSQ blocks",
      verbose=verbose,
      maxnthreads=maxnthreads )
   # save to file
   if dbpath != ""
      verbose && @info("Writing db to disk...")
      try
         flush(db)
      catch
         @warn("""something went wrong trying to save the db to disk, but the data
               should be ok; if it is crucial to keep it, try to save manually.""")
      end
      verbose && @info("... done")
   else
      verbose && @info("db is not written to disk since `dbpath` is empty.")
   end
   return db
end

function set_matrows!(d::Dat, okey::String, irows::Vector{Int})
   d.rows[okey] = irows
   return d
end

matrows(d::Dat, okey::String) = d.rows[okey]

function _alloc_lsq_matrix(configs, basis)
   # the indices associated with the basis are simply the indices within
   # the array - there is nothing else to do here
   #
   # loop through all observations and assign indices
   nrows = 0
   for (okey, d, _) in observations(configs)
      len = length(observation(d, okey))
      set_matrows!(d, okey, collect(nrows .+ (1:len)))
      nrows += len
   end
   # allocate and return the matrix
   return zeros(Float64, nrows, length(basis))
end

function safe_append!(db::LsqDB, db_lock, cfg, okey)
   # computing the lsq blocks ("rows") can be done in parallel,
   lsqrow = evallsq(cfg, basis(db), okey)
   irows = matrows(cfg, okey)
   # but writing them into the DB must be done in a threadsafe way
   lock(db_lock)
   db.Ψ[irows, :] .= lsqrow
   unlock(db_lock)
   return nothing
end

# ------------------- Evaluating LSQ Blocks -----------------

"""
Take a basis and split it into individual basis groups.
"""
function split_basis(basis; splitfun = b -> hash(Val{:basis}(), b))
   # get some basis hashs of the individual basis functions
   tps = splitfun.(basis)
   Iord = Vector{Int}[]
   Bord = Any[]
   for tp in unique(tps)
      # find which elements of basis have type `tp`
      I = findall( [tp == t  for t in tps] )
      push!(Iord, I)
      push!(Bord, [b for b in basis[I]])
   end
   return Bord, Iord
end

# fill the LSQ system, i.e. evaluate basis at data points
function evallsq(d::Dat, B::AbstractVector{TB}, okey) where {TB <: AbstractCalculator}
   B1 = [b for b in B]
   if !(isconcretetype(eltype(B1)))
      return evallsq_split(d, B1, okey)
   end
   # TB is a leaf-type so we can use "evaluate_many"
   return _evallsq(Val(Symbol(okey)), B1, Atoms(d))
end

"""
evaluate one specific kind of datum, such as energy, forces, etc
for one Dat across a large section of the basis; implicitly this will
normally be done via `evaluate_many` or similar. The line
```
vals = eval_obs(vDT, B, at)
```
should likely be the bottleneck of `evallsq`
"""
function _evallsq(vDT::Val,
                  B::AbstractVector{<:AbstractCalculator},
                  at::AbstractAtoms)
   # vals will be a vector containing multiple evaluations
   vals = eval_obs(vDT, B, at)
   # vectorise the first so we know the length of the observation
   vec1 = vec_obs(vDT, vals[1])
   # create a multi-D array to reshape these into
   A = Array{Float64}(undef, length(vec1), length(B))
   A[:, 1] = vec1
   for n = 2:length(B)
      A[:, n] = vec_obs(vDT, vals[n])
   end
   return A
end

"""
split the Basis `B` into subsets of identical types and evaluate
those independently (fast).
"""
function evallsq_split(d, basis, okey)
   # TB is not a leaf-type so we should split the basis to be able to
   # evaluate_many & co
   Bord, Iord = split_basis(basis)
   lenobs = length(d.D[okey])
   lsqrow = zeros(lenobs, length(basis))
   for (B, IB) in zip(Bord, Iord)
      lsqrow[:, IB] = evallsq(d, B, okey)
   end
   return lsqrow
end

# =========================== Convenience =================================

# TODO: THIS NEEDS TO BE REWRITTEN! => ALSO WRITE TESTS!
# function union(db1::LsqDB,db2::LsqDB; dbpath = (db1.dbpath * "_u"))
#    @info("Warning: the union implies that the data in the two databases are the same")
#    configtypes = collect(keys(db1.data_groups))
#    basis = cat(1, db1.basis, db2.basis)
#    data_groups = db1.data_groups
#
#    kron_groups1 = db1.kron_groups
#    kron_groups2 = db2.kron_groups
#
#    kron_groups = Dict{String, KronGroup}()
#    for k in configtypes
#       kron_groups[k] = KronGroup()
#       @assert haskey(db2.data_groups, k)
#       for ot in ["E", "F", "V"]
#          if !haskey(data_groups[k][1].D, ot)
#             continue
#          end
#          kron_groups[k][ot] = cat(3,kron_groups1[k][ot],kron_groups2[k][ot])
#       end
#    end
#    db = LsqDB(basis, data_groups, kron_groups, dbpath)
#    flush(db)
#    return db
# end

# function union(db1::LsqDB,db2::LsqDB, dbpath = (db1.dbpath * "_u"))
#    basis = [basis(db1), basis(db2)]
#    data_groups::Dict{String, DataGroup}
#       kron_groups::Dict{String, KronGroup}
#       dbpath::String
#    end
#
#    basis(db::LsqDB) = db.basis
#    basis(db::LsqDB, i::Integer) = db.basis[i]
#    data(db::LsqDB) = db.data
# end

configtypes(db::LsqDB) = unique(configtype.(db.configs))

_nconfigs(db::LsqDB, ct::AbstractString) =
   length(find(configtype.(db.configs) .== ct))

function info(db::LsqDB)
   # config names, how many
   all_cts = configtype.(db.configs)
   unq_cts = unique(all_cts)
   configs_info = Dict{String, Int}()
   for ct in unq_cts
      configs_info[ct] = sum(all_cts .== ct)
   end

   println("======================================================")
   println("       LsqDB Summary ")
   println("------------------------------------------------------")
   println(" Datagroup: configname  =>  number of configs ")
   for (i, (key, val)) in enumerate(configs_info)
      println("        $i : $key  =>  $val ")
   end
   # basis groups, how many, max degree
   B = db.basis
   Bord, Iord = split_basis(B)
   println("------------------------------------------------------")
   println(" Basis Group  | length | description ")
   for (i, B1) in enumerate(Bord)
      lenstr = replace(string(length(B1), pad=5), '0' => ' ')
      idxstr = replace(string(i, pad=2), '0' => ' ')
      desc = string(B1)
      desc = desc[1:min(50, length(desc))]
      println("          $idxstr  | $lenstr  | $desc" )
      #
   end
   println("======================================================")
end

end


module TestsUKBB

using Test
using BGEN
using CSV
using DataFrames
using TMLEEpistasis
using CategoricalArrays
using TOML

include("helper_fns.jl")

@testset "Test read_bgen function" begin
    # This file as an associated .bgi
    bgen_file = BGEN.datadir("example.8bits.bgen")
    b = TMLEEpistasis.read_bgen(bgen_file)
    @test b.idx isa BGEN.Index

    # This file as no associated .bgi
    bgen_file = BGEN.datadir("example.9bits.bgen")
    b = TMLEEpistasis.read_bgen(bgen_file)
    @test b.idx === nothing
end

@testset "Test samples_genotype" begin
    probabilities = [0.3 0.2 0.9;
                     0.5 0.2 0.05;
                     0.2 0.6 0.05]
    variant_genotypes = ["AA", "AT", "TT"]

    threshold = 0.9
    genotypes = TMLEEpistasis.samples_genotype(
        probabilities,
        variant_genotypes,
        threshold)
    @test genotypes[1] === genotypes[2] === missing
    @test genotypes[3] == "AA"

    threshold = 0.55
    genotypes = TMLEEpistasis.samples_genotype(
        probabilities,
        variant_genotypes,
        threshold)
    @test genotypes[1]  === missing
    @test genotypes[2] == "TT"
    @test genotypes[3] == "AA"

end

@testset "Test UKBBGenotypes function" begin
    build_query_file(threshold=0.95)
    queries = TMLEEpistasis.parse_queries(queryfile)
    genotypes = TMLEEpistasis.UKBBGenotypes(queryfile, queries)
    # I only look at the first 10 rows
    # SAMPLE_ID
    @test genotypes[1:9, "SAMPLE_ID"] == ["sample_00$i" for i in 1:9]
    # RSID_10
    @test genotypes[1:10, "RSID_10"] == repeat(["AG"], 10)
    # RSID_100
    @test all(genotypes[1:10, "RSID_100"] .=== ["AG", "AA", "AG", missing, "AG", "AG", missing, "AG", "GG", "AG"])

    rm(queryfile)
end

@testset "Test variant_genotypes" begin
    b = TMLEEpistasis.read_bgen(BGEN.datadir("example.8bits.bgen"))
    queries = [
        "query_1" => (RSID_10=["AA", "GG"], RSID_100=["AA", "GG"]),
        "query_2" => (RSID_10=["AA", "GG"], RSID_100=["AA", "GA"])
    ]
    # Defaults to "AG"
    v = variant_by_rsid(b, "RSID_10")
    @test TMLEEpistasis.variant_genotypes(v, queries) == ["AA", "AG", "GG"]
    # "GA" is specified by user so returned
    v = variant_by_rsid(b, "RSID_100")
    @test TMLEEpistasis.variant_genotypes(v, queries) == ["AA", "GA", "GG"]

end

@testset "Test preprocess" begin
    confounders = CSV.File(confoundersfile) |> DataFrame
    n = size(confounders)[1]
    genotypes = DataFrame(
        SAMPLE_ID = confounders.SAMPLE_ID,
        RSID_1    = vcat(Vector{Missing}(missing, 10), collect(1:n-10)),
        RSID_2    = vcat(collect(1:n-10), Vector{Missing}(missing, 10))
        )

    # Continuous phenotypes
    phenotypes = CSV.File(phenotypefile, select=["eid", "continuous_phenotype"]) |> DataFrame
    T, W, y = TMLEEpistasis.preprocess(genotypes, confounders, phenotypes;verbosity=0)

    @test y == phenotypes.continuous_phenotype[11:n-10]

    @test size(T) == (480, 2)
    @test size(W) == (480, 2)
    @test size(y) == (480,)

    @test names(T) == ["RSID_1", "RSID_2"]
    @test names(W) == ["PC1", "PC2"]

    @test T.RSID_1 == collect(1:n-20)
    @test T.RSID_1 isa CategoricalArray
    @test T.RSID_2 == collect(11:n-10)
    @test T.RSID_2 isa CategoricalArray

    # Binary phenotypes
    phenotypes = CSV.File(phenotypefile, select=["eid", "categorical_phenotype"]) |> DataFrame
    T, W, y = TMLEEpistasis.preprocess(genotypes, confounders, phenotypes;
                verbosity=0)

    @test y isa CategoricalArray{Bool,1,UInt32}
end

@testset "Test PhenotypeTMLEEpistasis with continuous target" begin
    # Only one continuous phenotype
    estimatorfile = joinpath("config", "tmle_config.toml")
    build_query_file()
    parsed_args = Dict(
        "phenotypes" => phenotypefile,
        "confounders" => confoundersfile,
        "queries" => queryfile,
        "estimator" => estimatorfile,
        "output" => "cont_results.csv",
        "phenotypes-list" => "data/phen_list_1.csv",
        "verbosity" => 0
    )

    UKBBVariantRun(parsed_args)

    results = CSV.File(parsed_args["output"]) |> DataFrame
    @test results.QUERYNAME == ["QUERY_1", "QUERY_2"]
    @test names(results) == ["PHENOTYPE", "QUERYNAME", "QUERYSTRING", "ESTIMATE", "PVALUE", "LOWER_BOUND", "UPPER_BOUND", "STD_ERROR", "QSTACK_COEFS", "NROWS", "TCOUNTS"]
    @test size(results) == (2, 11)

    # Clean
    rm(parsed_args["output"])
    rm(parsed_args["queries"])
end

@testset "Test PhenotypeTMLEEpistasis with categorical target and recovery mode" begin
    estimatorfile = joinpath("config", "tmle_config.toml")
    build_query_file()
    parsed_args = Dict(
        "phenotypes" => phenotypefile,
        "confounders" => confoundersfile,
        "queries" => queryfile,
        "estimator" => estimatorfile,
        "output" => "full_results.csv",
        "phenotypes-list" => "data/phen_list_2.csv",
        "verbosity" => 0
    )

    initial_results = DataFrame(PHENOTYPE=[:continuous_phenotype, :continuous_phenotype],
                    QUERYNAME=[:QUERY_DONE, :QUERY_DONE],
                    QUERYSTRING=["QUERYSTRING_1", "QUERYSTRING_2"],
                    ESTIMATE=[1., 2.],
                    PVALUE=[1., 2.],
                    LOWER_BOUND=[1., 2.],
                    UPPER_BOUND=[1., 2.],
                    STD_ERROR=[1., 2.],
                    QSTACK_COEFS=["TOTO", "TOTO"],
                    NROWS=[0, 0],
                    TCOUNTS=["TOTO", "TOTO"]
                    )
    CSV.write(parsed_args["output"], initial_results)

    UKBBVariantRun(parsed_args)

    results = CSV.File(parsed_args["output"]) |> DataFrame

    @test results[!, :NROWS] == [0, 0, 477, 477]
    @test results[!, :TCOUNTS] == ["TOTO", "TOTO", "RSID_10 & RSID_100: AG AA 65 | AG AG 248 | AG GG 145 | GG AA 7 | GG AG 4 | GG GG 8", "RSID_10 & RSID_100: AG AA 65 | AG AG 248 | AG GG 145 | GG AA 7 | GG AG 4 | GG GG 8"]
    @test results.QUERYNAME == ["QUERY_DONE", "QUERY_DONE", "QUERY_1", "QUERY_2"]
    @test names(results) == ["PHENOTYPE", "QUERYNAME", "QUERYSTRING", "ESTIMATE", "PVALUE", "LOWER_BOUND", "UPPER_BOUND", "STD_ERROR", "QSTACK_COEFS", "NROWS", "TCOUNTS"]
    @test size(results) == (4, 11)

    # Clean
    rm(parsed_args["output"])
    rm(parsed_args["queries"])
end

@testset "Test PhenotypeCrossValidation with categorical target and recovery mode" begin
    estimatorfile = joinpath("config", "tmle_config.toml")
    build_query_file()
    parsed_args = Dict(
        "phenotypes" => phenotypefile,
        "confounders" => confoundersfile,
        "queries" => queryfile,
        "estimator" => estimatorfile,
        "output" => "full_results.csv",
        "phenotypes-list" => "data/phen_list_2.csv",
        "verbosity" => 0
    )

    initial_results = DataFrame(PHENOTYPE=[:continuous_phenotype],
                    VARIANTS=["nothing"],
                    Q_RESULTSTRING=["nothing"],
                    G_RESULTSTRING=["nothing"],
                    )
    CSV.write(parsed_args["output"], initial_results)

    UKBBVariantRun(parsed_args, run_fn=TMLEEpistasis.PhenotypeCrossValidation)

    results = CSV.File(parsed_args["output"]) |> DataFrame

    @test results[2, "VARIANTS"] == "RSID_10 && RSID_100"
    # Check Q
    Qres = sort(split.(split(results[2, :Q_RESULTSTRING], " | "), ": "))
    @test Qres[1][1] == "HALClassifier_1"
    @test Qres[2][1] == "InteractionLMClassifier_1"
    @test Qres[3][1] == "XGBoostClassifier_1"
    for i in 1:3
        m_string, std_string = split(Qres[i][2], " ")
        parse(Float64, split(m_string, "=")[2])
        parse(Float64, split(std_string, "=")[2])
    end
    # Check G
    Gres = sort(split.(split(results[2, :G_RESULTSTRING], " | "), ": "))
    @test Gres[1][1] == "LogisticClassifier_1"
    @test Gres[2][1] == "XGBoostClassifier_1"
    for i in 1:2
        m_string, std_string = split(Gres[i][2], " ")
        parse(Float64, split(m_string, "=")[2])
        parse(Float64, split(std_string, "=")[2])
    end

    # Clean
    rm(parsed_args["output"])
    rm(parsed_args["queries"])
end

@testset "Test PhenotypeCrossValidation with continuous target" begin
    # Only one continuous phenotype
    estimatorfile = joinpath("config", "tmle_config.toml")
    build_query_file()
    parsed_args = Dict(
        "phenotypes" => phenotypefile,
        "confounders" => confoundersfile,
        "queries" => queryfile,
        "estimator" => estimatorfile,
        "output" => "cont_results.csv",
        "phenotypes-list" => "data/phen_list_1.csv",
        "verbosity" => 0
    )

    UKBBVariantRun(parsed_args, run_fn=TMLEEpistasis.PhenotypeCrossValidation)

    results = CSV.File(parsed_args["output"]) |> DataFrame

    @test results[1, "VARIANTS"] == "RSID_10 && RSID_100"
    # Check Q
    Qres = sort(split.(split(results[1, :Q_RESULTSTRING], " | "), ": "))
    @test Qres[1][1] == "HALRegressor_1"
    @test Qres[2][1] == "InteractionLMRegressor_1"
    @test Qres[3][1] == "XGBoostRegressor_1"
    @test Qres[4][1] == "XGBoostRegressor_2"
    for i in 1:4
        m_string, std_string = split(Qres[i][2], " ")
        parse(Float64, split(m_string, "=")[2])
        parse(Float64, split(std_string, "=")[2])
    end
    # Check G
    Gres = sort(split.(split(results[1, :G_RESULTSTRING], " | "), ": "))
    @test Gres[1][1] == "LogisticClassifier_1"
    @test Gres[2][1] == "XGBoostClassifier_1"
    for i in 1:2
        m_string, std_string = split(Gres[i][2], " ")
        parse(Float64, split(m_string, "=")[2])
        parse(Float64, split(std_string, "=")[2])
    end

    # Clean
    rm(parsed_args["output"])
    rm(parsed_args["queries"])
end

end;

true


module Strided

import Base: parent, size, strides, tail, setindex
using Base: @propagate_inbounds, RangeIndex, Dims
const SliceIndex = Union{RangeIndex,Colon}

using LinearAlgebra

using TupleTools
using TupleTools: StaticLength

export StridedView, @strided, @unsafe_strided, sreshape, sview

# function __init__()
#     LinearAlgebra.BLAS.set_num_threads(1)
#     Threads.nthreads() == 1 && warn("Strided disables BLAS multithreading, enable Julia threading (`export JULIA_NUM_THREADS = N`) to benefit from multithreaded matrix multiplication and more")
# end

const _use_threads = Ref(true)
use_threads() = _use_threads[]

"""
    disable_threads()

Disable the use of threading in Strided.
Also see [`enable_threads()`](@ref)
"""
function disable_threads()
    _use_threads[] = false
    return
end

"""
    enable_threads()

(Re)-enable threading in Strided.
"""
function enable_threads()
    _use_threads[] = true
    return
end

# used to factor the number of threads
function simpleprimefactorization(n::Int)
    factors = Vector{Int}()
    k = 2
    while k <= n
        d, r = divrem(n, k)
        if r == 0
            push!(factors, k)
            n = d
        else
            k += 1
        end
    end
    return factors
end

const factors = Vector{Int}(undef, 0)

function __init__()
    f = reverse!(simpleprimefactorization(Threads.nthreads()))
    resize!(factors, length(f))
    copyto!(factors, f)
end

include("abstractstridedview.jl")
include("stridedview.jl")
include("unsafestridedview.jl")
include("mapreduce.jl")
include("broadcast.jl")
include("macros.jl")

# function _precompile_()
#     ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
#     for T in (Float64, ComplexF64)
#         USMatrix = UnsafeStridedView{T, 2, T, typeof(Base.identity)}
#         SMatrix = StridedView{T, 2, Vector{T}, typeof(Base.identity)}
#
#         precompile(Tuple{typeof(mul!), USMatrix, USMatrix, USMatrix})
#         precompile(Tuple{typeof(mul!), USMatrix, USMatrix, USMatrix, T, T})
#
#         precompile(Tuple{typeof(mul!), SMatrix, SMatrix, SMatrix})
#         precompile(Tuple{typeof(mul!), SMatrix, SMatrix, SMatrix, T, T})
#
#         if T <: Complex
#             CUSMatrix = UnsafeStridedView{T, 2, T, typeof(Base.conj)}
#             CSMatrix = StridedView{T, 2, Vector{T}, typeof(Base.conj)}
#
#             precompile(Tuple{typeof(mul!), USMatrix, USMatrix, CUSMatrix})
#             precompile(Tuple{typeof(mul!), USMatrix, CUSMatrix, USMatrix})
#             precompile(Tuple{typeof(mul!), USMatrix, CUSMatrix, CUSMatrix})
#
#             precompile(Tuple{typeof(mul!), SMatrix, SMatrix, CSMatrix})
#             precompile(Tuple{typeof(mul!), SMatrix, CSMatrix, SMatrix})
#             precompile(Tuple{typeof(mul!), SMatrix, CSMatrix, CSMatrix})
#         end
#
#         for N = 1:4
#             USArray = UnsafeStridedView{T, N, T, typeof(Base.identity)}
#             SArray = StridedView{T, N, Vector{T}, typeof(Base.identity)}
#
#             for ASArray in (USArray, SArray)
#                 precompile(Tuple{typeof(copyto!), ASArray, ASArray})
#                 precompile(Tuple{typeof(mul!), ASArray, T, ASArray})
#                 precompile(Tuple{typeof(mul!), ASArray, ASArray, T})
#                 precompile(Tuple{typeof(rmul!), ASArray, T})
#                 precompile(Tuple{typeof(lmul!), T, ASArray})
#                 precompile(Tuple{typeof(axpy!), T, ASArray, ASArray})
#                 precompile(Tuple{typeof(axpby!), T, ASArray, T, ASArray})
#             end
#
#             if T <: Complex
#                 CUSArray = UnsafeStridedView{T, N, T, typeof(Base.conj)}
#                 CSArray = StridedView{T, N, Vector{T}, typeof(Base.conj)}
#
#                 for (ASArray, CASArray) in ((USArray, CUSArray), (SArray, CSArray))
#                     precompile(Tuple{typeof(copyto!), ASArray, CASArray})
#                     precompile(Tuple{typeof(mul!), ASArray, T, CASArray})
#                     precompile(Tuple{typeof(mul!), ASArray, CASArray, T})
#                     precompile(Tuple{typeof(rmul!), CASArray, T})
#                     precompile(Tuple{typeof(lmul!), T, CASArray})
#                     precompile(Tuple{typeof(axpy!), T, CASArray, ASArray})
#                     precompile(Tuple{typeof(axpby!), T, CASArray, T, ASArray})
#                 end
#             end
#         end
#     end
# end
# _precompile_()

end


function get_RealQuantity(xmle::LightXML.XMLElement, str::String)
  dx = 0.0
  x = 0.0
  for z in child_elements(xmle)
    if name(z) == str
      for y in child_elements(z)
        if name(y) == "value"
          x = parse(Float64, content(y))
        elseif name(y) == "uncertainty"
          dx = parse(Float64, content(y))
        end
      end
    end
  end
  return (x, dx)
end

function parse_qml(evt::XMLElement)
  MT = Array{SeisSrc,1}(undef, 0)
  evt_src = attribute(evt, "publicID")
  evt_id = split(evt_src, r"[;:/=]")[end]
  preferredOriginID = ""
  preferredMagnitudeID = ""
  preferredFocalMechanismID = ""

  mags = Array{XMLElement,1}(undef,0)
  origs = Array{XMLElement,1}(undef,0)
  mechs = Array{XMLElement,1}(undef,0)
  eqtype = ""

  for ch in child_elements(evt)
    cn = name(ch)
    if cn == "preferredOriginID"
      preferredOriginID = split(content(ch), r"[;:/=]")[end]
    elseif cn == "preferredMagnitudeID"
      preferredMagnitudeID = split(content(ch), r"[;:/=]")[end]
    elseif cn == "preferredFocalMechanismID"
      preferredFocalMechanismID = split(content(ch), r"[;:/=]")[end]
    elseif cn == "type"
      eqtype = content(ch)
    elseif cn == "magnitude"
      push!(mags, ch)
    elseif cn == "origin"
      push!(origs, ch)
    elseif cn == "focalMechanism"
      push!(mechs, ch)
    end
  end

  # ==========================================================================
  # Focal Mechanism
  n = 0
  while n < length(mechs)
    n         = n+1
    auth      = ""
    mech      = getindex(mechs, n)
    np        = zeros(Float64, 3, 2)
    pax       = zeros(Float64, 3, 3)
    mt        = zeros(Float64, 6)
    dm        = fill!(zeros(Float64, 6), Inf)

    # generate SeisSrc object
    S = SeisSrc()
    ssrc = attribute(mech, "publicID")
    setfield!(S, :id, String(split(ssrc, r"[;:/=]")[end]))

    for child in child_elements(mech)
      cname = name(child)

      if cname == "creationInfo"
        for grandchild in child_elements(child)
          if name(grandchild) ==  "author"
            auth = content(grandchild)
          end
        end

      elseif cname == "azimuthalGap"
        setfield!(S, :gap, parse(Float64, content(child)))

      elseif cname == "stationPolarityCount"
        setfield!(S, :npol, parse(Int64, content(child)))

      elseif cname == "nodalPlanes"
        j = 0
        for np2 in child_elements(child)  # nodalPlane1, nodalPlane2
          if name(np2) != "preferredPlane"
            for grandchild in child_elements(np2) # strike, dip, rake
              for greatgrandchild in child_elements(grandchild)
                if name(greatgrandchild) == "value"
                  j = j+1
                  setindex!(np, parse(Float64, content(greatgrandchild)), j)
                end
              end
            end
          end
        end
        S.misc["planes_desc"] = "strike, dip, rake"

      elseif cname == "principalAxes"
        j = 0
        for grandchild in child_elements(child) # tAxis, pAxis, nAxis
          for greatgrandchild in child_elements(grandchild) # azimuth, plunge, length
            for redheadedstepchild in child_elements(greatgrandchild)
              if name(redheadedstepchild) == "value"
                j = j+1
                setindex!(pax, parse(Float64, content(greatgrandchild)), j)
              end
            end
          end
        end
        S.misc["pax_desc"] = "azimuth, plunge, length"

      elseif cname == "momentTensor"
        S.misc["mt_id"] = attribute(child, "publicID")
        for a in child_elements(child)
          aname = name(a)
          if aname == "scalarMoment"
            setfield!(S, :m0, parse(Float64, content(a)))
          elseif aname == "tensor"
            j = 0
            for k in ("Mrr", "Mtt", "Mpp", "Mrt", "Mrp", "Mtp")
              j = j+1
              (mt[j], dm[j]) = get_RealQuantity(a, k)
            end
          elseif aname == "sourceTimeFunction"
            for b in child_elements(a)
              if name(b) == "type"
                S.st.desc = content(b)
              elseif name(b) == "duration"
                S.st.dur = parse(Float64, content(b))
              elseif name(b) == "riseTime"
                S.st.rise = parse(Float64, content(b))
              elseif name(b) == "decayTime"
                S.st.decay = parse(Float64, content(b))
              end
            end
          elseif aname == "derivedOriginID"
            S.misc["derivedOriginID"] = String(split(content(a), r"[;:/=]")[end])
            S.misc["xmt_derivedOriginID"] = string(a)
          else
            S.misc["xmt_" * aname] = string(a)
          end
        end
      elseif cname == "methodID"
        S.misc["methodID"] = content(child)
      else
        S.misc["xmech_" * cname] = string(child)
      end

    end
    src_str = ssrc * "," * auth
    setfield!(S, :planes, np)
    setfield!(S, :pax, pax)
    setfield!(S, :mt, mt)
    setfield!(S, :dm, dm)
    setfield!(S, :src, src_str)
    note!(S, "+origin ¦ " * src_str)
    push!(MT, S)
  end

  # ==========================================================================
  # Choose a focal mechanism to retain; then determine mag_id and loc_id
  mag_id = ""
  loc_id = ""
  if length(MT) > 0
    # Store moment tensor with lowest M_err
    if isempty(preferredFocalMechanismID)
      sort!(MT, by=x->sum(abs.(x.dm)))
      R = getindex(MT, 1)
    else
      R = getindex(MT, findfirst([occursin(getfield(m, :id), preferredFocalMechanismID) for m in MT]))
    end

    # Set ID
    if isempty(preferredMagnitudeID)
      mag_id = get(getfield(R, :misc), "derivedOriginID", "")
    else
      mag_id = preferredMagnitudeID
    end
    loc_id = isempty(preferredOriginID) ? mag_id : preferredOriginID
  else
    R = SeisSrc()
  end

  # ==========================================================================
  # Magnitude
  H = SeisHdr()
  setfield!(R, :eid, identity(getfield(H, :id)))

  n = 0
  oid = ""
  pid = ""
  auth = ""
  m = -5.0f0
  msc = ""
  nst = 0
  gap = 0.0
  while n < length(mags)
    n = n+1
    mag = getindex(mags, n)
    oid = ""
    pid = ""
    auth = ""
    m = -5.0f0
    msc = ""
    nst = 0
    gap = 0.0

    # src
    originID = get_elements_by_tagname(mag, "originID")
    pid = attribute(mag, "publicID")
    oid = isempty(originID) ? "" : content(first(originID))

    # msc
    mtype = get_elements_by_tagname(mag, "type")
    msc = isempty(mtype) ? "" : content(first(mtype))

    if occursin(oid, mag_id) || n == 1 || startswith(lowercase(msc), "mw")
      for child in child_elements(mag)
        cname = name(child)
        if cname == "mag"
          m = parse(Float32, content(first(get_elements_by_tagname(child, "value"))))
        elseif cname == "stationCount"
          nst = parse(Int64, content(child))
        elseif cname == "azimuthalGap"
          gap = parse(Float64, content(child))
        elseif cname == "creationInfo"
          for grandchild in child_elements(child)
            if name(grandchild) ==  "author"
              auth = content(grandchild)
            end
          end
        end
      end
    end
  end
  MAG = EQMag(m, msc, nst, gap, pid * "," * oid * "," * auth)
  if MAG.src == ",,"
    MAG.src = ""
  end

  # ==========================================================================
  # Location
  n = 0
  nst = 0
  ot = "1970-01-01T00:00:00"
  gap = 0.0
  auth = ""
  ltyp = ""
  loc_src = ""
  locflags = Array{Char, 1}(undef,8)
  fill!(locflags, '0')
  loc = zeros(Float64, 12)
  while n < length(origs)
    n = n+1
    orig = getindex(origs, n)

    # _______________________________________________________
    # Only parse locations corresponding to a desirable ID
    if occursin(attribute(orig, "publicID"), loc_id) || (n == 1)
      loc_src = attribute(orig, "publicID")

      # Reset temp variables
      nst = 0
      gap = 0.0
      auth = ""
      ltyp = ""
      ot = "1970-01-01T00:00:00"
      fill!(locflags, '0')
      fill!(loc, zero(Float64))

      # Try to set location first
      j = 0
      for str in ("latitude", "longitude", "depth")
        j = j + 1
        (loc[j], dloc) = get_RealQuantity(orig, str)
        if j == 1
          loc[5] = dloc
        elseif j == 2
          loc[4] = dloc
        elseif j == 3
          setindex!(loc, getindex(loc, 3)*0.001, 3)
          loc[6] = dloc
        end
      end

      # Now loop
      for child in child_elements(orig)
        cname = name(child)

        #= originUncertainty removed 2019-11-01
          No documentation of originUncertainty is known to exist;
          can't ascertain intended meaning/use of originUncertainty;
          can't tell if observatories use it uniformly
        =#

        if cname == "type"
          ltyp_tmp = content(child)
          if ltyp_tmp != ""
            ltyp = ltyp_tmp
          end

        elseif cname == "time"
          for grandchild in child_elements(child)
            gcname = name(grandchild)
            if gcname == "value"
              ot = content(grandchild)
            elseif gcname == "uncertainty"
              loc[7] = parse(Float64, content(grandchild))
            end
          end

        elseif cname == "quality"
          for grandchild in child_elements(child)
            gcname = name(grandchild)
            if gcname == "standardError"
              loc[8] = parse(Float64, content(grandchild))
            elseif gcname == "azimuthalGap"
              loc[10] = parse(Float64, content(grandchild))
            elseif gcname == "minimumDistance"
              loc[11] = parse(Float64, content(grandchild))
            elseif gcname == "maximumDistance"
              loc[12] = parse(Float64, content(grandchild))
            elseif gcname == "associatedStationCount"
              nst = parse(Int64, content(grandchild))
            end
          end

        elseif cname == "epicenterFixed"
          locflags[1] = content(child)[1]
          locflags[2] = locflags[1]

        elseif cname == "depthType"
          if content(child) == "operator assigned"
            locflags[3] = '1'
          end

        elseif cname == "timeFixed"
          locflags[4] = content(child)[1]

        elseif cname == "creationInfo"
          for grandchild in child_elements(child)
            if name(grandchild) ==  "author"
              auth = content(grandchild)
            end
          end
        elseif (cname in ("latitude", "longitude", "depth")) == false
          H.misc["xloc_" * cname] = string(child)
        end
      end
      # _______________________________________________________

    end
  end
  LOC = EQLoc(loc..., nst, parse(UInt8, join(locflags), base=2), "", ltyp, "", loc_src * "," * auth )

  setfield!(H, :id, String(evt_id))
  setfield!(H, :loc, LOC)
  setfield!(H, :ot, DateTime(replace(ot, r"[A-S,U-Z,a-z]" => "")[1:min(end,23)]))
  setfield!(H, :mag, MAG)
  setfield!(H, :typ, eqtype)
  setfield!(H, :src, evt_src)
  setfield!(R, :eid, String(evt_id))
  note!(H, "+origin ¦ " * evt_src)

  # Post-process: ensure these are empty if not initialized
  if R.mt == zeros(Float64,6)
    R.mt = Float64[]
  end

  if R.dm == Inf64.*ones(Float64,6)
    R.dm = Float64[]
  end

  if R.pax == zeros(Float64, 3, 3)
    R.pax = Array{Float64, 2}(undef, 0, 0)
  end

  if R.planes == zeros(Float64, 3, 2)
    R.planes = Array{Float64, 2}(undef, 0, 0)
  end

  if R.src == ","
    R.src = ""
  end
  return H, R
end

function event_xml!(EvCat::Array{SeisHdr,1}, EvSrc::Array{SeisSrc, 1}, xdoc::XMLDocument)
  qxml  = first(child_elements(root(xdoc)))
  elts  = child_elements(qxml)
  for elt in elts
    if name(elt) == "event"
      H, R = parse_qml(elt)
      push!(EvCat, H)
      push!(EvSrc, R)

    end
  end
  return nothing
end

"""
    EvCat, EvSrc = read_qml(fpat::String)

Read QuakeML files matching string pattern `fpat`. Returns an array of `SeisHdr`
objects as `EvCat` and an array of `SeisSrc` objects as `EvSrc`, such that
`EvCat[i]` and `EvSrc[i]` describe the preferred location (origin) and
preferred event source (focal mechanism or moment tensor) of event `i`.

"""
function read_qml(fpat::String)
  files = safe_isfile(fpat) ? [fpat] : ls(fpat)
  EvCat = Array{SeisHdr,1}()
  EvSrc = Array{SeisSrc,1}()
  for file in files
    xdoc  = parse_file(file)
    k = length(EvCat)
    event_xml!(EvCat, EvSrc, xdoc)
    L = length(EvCat)
    fstr = abspath(file)
    for i in k+1:L
      setfield!(EvCat[i], :src, fstr)
      note!(EvCat[i], "+source ¦ " * fstr)
      setfield!(EvSrc[i], :src, fstr)
      note!(EvSrc[i], "+source ¦ " * fstr)
    end
    free(xdoc)
  end
  return EvCat, EvSrc
end

function is_qml(k::String, s::String)
  (startswith(s, "<") && endswith(s, ">")) || return false
  i1 = first(findnext("<", s,1))
  j1 = last(findnext(">", s,1))
  s1 = s[nextind(s, i1):prevind(s, j1)]
  (split(k, "_")[2] == s1) || return false
  i2 = first(findlast("<", s))
  j2 = last(findlast(">", s))
  s2 = s[nextind(s, i2):prevind(s, j2)]
  ("/" * s1) == s2 || return false
  return true
end

function new_qml!(io::IO)
  write(io, "<?xml version=\"1.0\" encoding=\"UTF-8\"?>
  <quakeml xmlns=\"http://quakeml.org/xmlns/quakeml/1.2\">\n  <eventParameters publicID=\"smi:SeisIO.jl\">\n    <creationInfo>\n      <agencyID>SeisIO</agencyID>\n      <creationTime>")
  print(io, now())
  write(io, "</creationTime>\n    </creationInfo>\n")
  return nothing
end

function write_misc(io::IO, D::Dict{String,Any}, pref::String, p::Int64)
  for k in keys(D)
    if startswith(k, pref)
      xv = D[k]
      if is_qml(k, xv)
        write(io, " "^2p)
        write(io, xv)
        write(io, "\n")
      end
    end
  end
  return nothing
end

function write_real(io::IO, str::String, x::Union{DateTime,AbstractFloat}, p::Int64)
  write(io, " "^2p, "<")
  write(io, str)
  write(io, ">\n", " "^(2p+2), "<value>")
  print(io, x)
  write(io, "</value>\n", " "^2p, "</")
  write(io, str)
  write(io, ">\n")
  return nothing
end

function write_real(io::IO, str::String, x::Union{DateTime,AbstractFloat}, dx::AbstractFloat, p::Int64)
  write(io, " "^2p, "<")
  write(io, str)
  write(io, ">\n", " "^(2p+2), "<value>")
  print(io, x)
  write(io, "</value>\n", " "^(2p+2), "<uncertainty>")
  print(io, dx)
  write(io, "</uncertainty>\n", " "^2p, "</")
  write(io, str)
  write(io, ">\n")
  return nothing
end

function write_pax(io::IO, pax::Array{Float64,2})
  str = ("t", "p", "n")
  nP = size(pax, 2)
  write(io, "        <principalAxes>\n")
  for i in 1:nP
    write(io, "          <", str[i], "Axis>\n")
    write(io, "            <azimuth>\n              <value>")
    print(io, pax[1,i])
    write(io, "</value>\n            </azimuth>\n")
    write(io, "            <plunge>\n              <value>")
    print(io, pax[2,i])
    write(io, "</value>\n            </plunge>\n")
    write(io, "            <length>\n              <value>")
    print(io, pax[3,i])
    write(io, "</value>\n            </length>\n")
    write(io, "          </", str[i], "Axis>\n")
  end
  write(io, "        </principalAxes>\n")
  return nothing
end

function write_qml!(io::IO, HDR::Array{SeisHdr,1}, SRC::Array{SeisSrc,1}, v::Integer)
  Ri = zeros(Int64, length(HDR))
  R_id = Array{String,1}(undef, length(SRC))
  for i in 1:length(SRC)
    R_id[i] = getfield(SRC[i], :eid)
  end

  for i in 1:length(HDR)
    id = getfield(HDR[i], :id)
    for j in 1:length(R_id)
      if id == R_id[j]
        Ri[i] = j
        break
      end
    end
  end

  for i in 1:length(HDR)
    H = HDR[i]
    (v > 0) && println("Writing event ", H.id)

    i_id = findfirst([occursin("origin", n) for n in H.notes])
    id_str = String(split(i_id == nothing ? " ¦  ¦ " : H.notes[i_id], " ¦ ")[end])
    write(io, "    <event publicID=\"")
    print(io, id_str)
    write(io, "\">\n")
    if isempty(H.loc)
      (v > 0) && println("  Skipped location (H.loc empty)")
    else
      loc_orig, loc_auth, xx = split_id(H.loc.src, c=",")
      write(io, "      <preferredOriginID>")
      write(io, H.id)
      write(io, "</preferredOriginID>\n      <type>")
      write(io, H.typ)
      write(io, "</type>\n")

      # ---------------------------------------------------
      # Origin
      write(io, "      <origin publicID=\"", loc_orig, "\">\n")
      if H.loc.dt > 0.0
        write_real(io, "time", H.ot, H.loc.dt, 4)
      else
        write_real(io, "time", H.ot, 4)
      end
      L = H.loc

      # lat, lon, dep
      if L.dy == 0.0
        write_real(io, "latitude", L.lat, 4)
      else
        write_real(io, "latitude", L.lat, L.dy, 4)
      end
      if L.dx == 0.0
        write_real(io, "longitude", L.lon, 4)
      else
        write_real(io, "longitude", L.lon, L.dx, 4)
      end
      if L.dz == 0.0
        write_real(io, "depth", L.dep*1000.0, 4)
      else
        write_real(io, "depth", L.dep*1000.0, L.dz, 4)
      end

      # flags
      flags = falses(4)
      for n = 1:4
        flags[n] = >>(<<(H.loc.flags, n-1),7)
      end
      if flags[3]
        write(io, "        <depthType>operator assigned</depthType>\n")
      end
      if flags[4]
        write(io, "        <timeFixed>1</timeFixed>\n")
      end
      if flags[1] || flags[2]
        write(io, "        <epicenterFixed>1</epicenterFixed>\n")
      end

      # Location quality
      do_qual = false
      for f in loc_qual_fields
        if getfield(L, f) != 0.0
          do_qual = true
          break
        end
      end
      if do_qual || (L.nst > 0)
        write(io, "        <quality>\n")

        if L.nst > 0
          write(io, "          <associatedStationCount>")
          print(io, L.nst)
          write(io, "</associatedStationCount>\n")
        end

        for (i,f) in enumerate(loc_qual_fields)
          j = getfield(L, f)
          if j != 0.0
            write(io, "          <")
            write(io, loc_qual_names[i])
            write(io, ">")
            print(io, j)
            write(io, "</")
            write(io, loc_qual_names[i])
            write(io, ">\n")
          end
        end
        write(io, "        </quality>\n")
      end

      if !isempty(L.typ)
        write(io, "        <type>")
        print(io, L.typ)
        write(io, "</type>\n")
      end

      # Author
      if !isempty(loc_auth)
        write(io, "        <creationInfo>\n          <author>")
        print(io, loc_auth)
        write(io, "</author>\n        </creationInfo>\n")
      end

      # other location properties
      write_misc(io, H.misc, "xloc_", 4)

      # done Origin
      write(io, "      </origin>\n")
    end

    # ---------------------------------------------------
    # Focal Mechanism
    if Ri[i] == 0
      (v > 0) && println("  Skipped focal mechanism (no R.eid matches)\n  Skipped moment tensor (fields empty)")
    else
      j = Ri[i]
      R = SRC[j]

      i_src = findfirst([occursin("origin", n) for n in R.notes])
      src_str = String(split(i_src == nothing ? " ¦  ¦ " : R.notes[i_src], " ¦ ")[end])
      foc_orig, foc_auth, xx = split_id(src_str, c=",")
      write(io, "      <focalMechanism publicID=\"")
      write(io, foc_orig)
      write(io, "\">\n")

      # Nodal planes
      if !isempty(R.planes)
        write(io, "        <nodalPlanes>\n")
        nP = size(R.planes, 2)
        for i in 1:nP
          write(io, "          <nodalPlane")
          print(io, i)
          write(io, ">\n")
          write(io, "            <strike>\n              <value>")
          print(io, R.planes[1,i])
          write(io, "</value>\n            </strike>\n")
          write(io, "            <dip>\n              <value>")
          print(io, R.planes[2,i])
          write(io, "</value>\n            </dip>\n")
          write(io, "            <rake>\n              <value>")
          print(io, R.planes[3,i])
          write(io, "</value>\n            </rake>\n")
          write(io, "          </nodalPlane")
          print(io, i)
          write(io, ">\n")
        end
        write(io, "        </nodalPlanes>\n")
      end

      # Principal axes
      if !isempty(R.pax)
        write_pax(io, R.pax)
      end

      # Azimuthal gap
      if R.gap != 0.0
        write(io, "        <azimuthalGap>")
        print(io, R.gap)
        write(io, "</azimuthalGap>\n")
      end

      # Polarity count
      if R.npol != 0
        write(io, "        <stationPolarityCount>")
        print(io, R.npol)
        write(io, "</stationPolarityCount>\n")
      end

      # methodID
      if haskey(R.misc, "methodID")
        write(io, "        <methodID>")
        write(io, get(R.misc, "methodID", ""))
        write(io, "</methodID>\n")
      end

      # Author
      if !isempty(foc_auth)
        write(io, "        <creationInfo>\n            <author>")
        write(io, foc_auth)
        write(io, "</author>\n          </creationInfo>\n")
      end

      # Moment Tensor
      if (isempty(R.mt) && (R.m0 == 0.0) && isempty(R.st))
        (v > 0) && println("  Skipped moment tensor (fields empty)")
      else
        write(io, "        <momentTensor publicID=\"")
        mt_id = haskey(R.misc, "mt_id") ? R.misc["mt_id"] : "smi:SeisIO/moment_tensor;fmid=" * R.id
        write(io, mt_id)
        write(io, "\">\n")

        if R.m0 != 0.0
          write_real(io, "scalarMoment", R.m0, 5)
        end

        if isempty(R.mt) == false
          write(io, "          <tensor>\n")
          mt_strings = ("Mrr", "Mtt", "Mpp", "Mrt", "Mrp", "Mtp")
          for i = 1:length(R.mt)
            write_real(io, mt_strings[i], R.mt[i], R.dm[i], 6)
          end
          write(io, "          </tensor>\n")
          write_misc(io, R.misc, "xmt_", 5)
        end

        if !isempty(R.st)
          write(io, "          <sourceTimeFunction>\n")
          write(io, "            <type>")
          write(io, R.st.desc)
          write(io, "</type>\n")
          write(io, "            <duration>")
          print(io, R.st.dur)
          write(io, "</duration>\n")
          write(io, "            <riseTime>")
          print(io, R.st.rise)
          write(io, "</riseTime>\n")
          write(io, "            <decayTime>")
          print(io, R.st.decay)
          write(io, "</decayTime>\n")
          write(io, "          </sourceTimeFunction>\n")
        end
        write(io, "        </momentTensor>\n")
      end

      write(io, "      </focalMechanism>\n")
    end

    # ---------------------------------------------------
    # Magnitude
    if isempty(H.mag)
      (v > 0) && println("  Skipped magnitude (H.mag empty)")
    else
      mag_pid, mag_orig, mag_auth, xx = split_id(H.mag.src, c=",")
      write(io, "      <magnitude publicID=\"")
      write(io, mag_pid)
      write(io, "\">\n")
      write_real(io, "mag", H.mag.val, 4)
      write(io, "        <type>", H.mag.scale, "</type>\n")
      isempty(mag_orig) || write(io, "        <originID>", mag_orig, "</originID>\n")

      if !isempty(mag_auth)
        write(io, "        <creationInfo>\n          <author>", mag_auth, "</author>\n        </creationInfo>\n")
      end

      if H.mag.gap != 0.0
        write(io, "        <azimuthalGap>")
        print(io, H.mag.gap)
        write(io, "</azimuthalGap>\n")
      end

      if H.mag.nst > 0
        write(io, "        <stationCount>")
        print(io, H.mag.nst)
        write(io, "</stationCount>\n")
      end

      write(io, "      </magnitude>\n")
    end
    write(io, "    </event>\n")
  end
  write(io, "</eventParameters>\n</quakeml>\n")
  return nothing
end

@doc """
    write_qml(fname, Ev::SeisEvent; v::Integer=0)

Write event metadata from SeisEvent `Ev` to file `fname`.

    write_qml(fname, SHDR::SeisHdr; v::Integer=0)
    write_qml(fname, SHDR::Array{SeisHdr,1}; v::Integer=0)

Write QML to file `fname` from `SHDR`.

If `fname` exists, and is QuakeML, SeisIO appends the existing XML. If the
file exists, but is NOT QuakeML, an error is thrown; the file isn't overwritten.

    write_qml(fname, SHDR::SeisHdr, SSRC::SeisSrc; v::Integer=0)
    write_qml(fname, SHDR::Array{SeisHdr,1}, SSRC::Array{SeisSrc,1}; v::Integer=0)

Write QML to file `fname` from `SHDR` and `SSRC`.

!!! warning

    To write data from `R ∈ SSRC`, it must be true that `R.eid == H.id` for some `H ∈ SHDR`.
""" write_qml
function write_qml(fname::String, HDR::Array{SeisHdr,1}, SRC::Array{SeisSrc,1}; v::Integer=0)
  H0 = SeisHdr[]
  R0 = SeisSrc[]

  if safe_isfile(fname)
    io = open(fname, "a+")

    # test whether file can be appended
    fastseekend(io)
    fastskip(io, -30)
    test_str = String(read(io))
    if test_str == "</eventParameters>\n</quakeml>\n"

      # behavior for files that are QuakeXML as produced by SeisIO
      fastskip(io, -30)
    else
      try

        # file exists and is readable QuakeXML but not produced by SeisIO
        seekstart(io)
        fstart = String(fastread(io, 5))
        if fstart == "<?xml"
          close(io)
          append!(H0, HDR)
          append!(R0, SRC)
          (H1, R1) = read_qml(fname)
          @assert (isempty(H1) == false)
          @assert (isempty(R1) == false)
          append!(H0, H1)
          append!(R0, R1)
        else
          error("incompatible file type!")
        end
      catch err

        # file exists but isn't QuakeXML
        @warn(string(fname, " isn't valid QuakeXML; can't append, exit with error!"))
        rethrow(err)
      end

      io = open(fname, "w")
      new_qml!(io)
    end
  else

    # new file
    io = open(fname, "w")
    new_qml!(io)
  end

  opts = string(", v=", v)
  if isempty(H0) && isempty(R0)
    write_qml!(io, HDR, SRC, v)
    for h in HDR
      fwrite_note_quake!(h, "asdf_wqml", fname, opts)
    end
    for r in SRC
      fwrite_note_quake!(r, "asdf_wqml", fname, opts)
    end
  else
    write_qml!(io, H0, R0, v)
    for h in H0
      fwrite_note_quake!(h, "asdf_wqml", fname, opts)
    end
    for r in R0
      fwrite_note_quake!(r, "asdf_wqml", fname, opts)
    end
  end
  close(io)
  return nothing
end
write_qml(fname::String, H::SeisHdr, R::SeisSrc; v::Integer=0) = write_qml(fname, [H], [R], v=v)
write_qml(fname::String, HDR::Array{SeisHdr,1}; v::Integer=0) = write_qml(fname, HDR, SeisSrc[], v=v)
write_qml(fname::String, H::SeisHdr; v::Integer=0) = write_qml(fname, [H], SeisSrc[], v=v)
write_qml(fname::String, Ev::SeisEvent; v::Integer=0) = write_qml(fname, [Ev.hdr], [Ev.source], v=v)


module GateMap
# =========================================================================== #
#                           Imported modules
# =========================================================================== #

using Base
using LinearAlgebra
using StaticArrays
import QXZoo.GateOps
import QXZoo.DefaultGates

# =========================================================================== #
#                           Exported functions
# =========================================================================== #

export create_gate_1q, create_gate_2q

# =========================================================================== #
#                       Gate caching and retrieval
# =========================================================================== #

"""Module gate cache

Since many different circuits may use the same gates, keeping a module-global
cache makes sense to avoid recreating them. Each subcircuit can have a subset of
the global cache's gates."""
gates = Dict{Union{GateOps.GateSymbol, GateOps.GateSymbolP, Function}, Function}()

# TODO
# const gates = let .... end

"""
    init_cache()

Initialise the cache for default gates used in circuit generation. Modules generating
gates not stored here should cache them for use by other circuits.

"""
function init_cache() #gates::Dict{Union{<:GateOps.AGateSymbol, Function}, Function})
    push!(gates, DefaultGates.GateSymbols.p00 => p00)
    push!(gates, DefaultGates.GateSymbols.p10 => p10)
    push!(gates, DefaultGates.GateSymbols.p01 => p01)
    push!(gates, DefaultGates.GateSymbols.p11 => p11)
    push!(gates, DefaultGates.GateSymbols.I => I)
    push!(gates, DefaultGates.GateSymbols.h => h)
    push!(gates, DefaultGates.GateSymbols.x => x)
    push!(gates, DefaultGates.GateSymbols.y => y)
    push!(gates, DefaultGates.GateSymbols.z => z)

    push!(gates, DefaultGates.GateSymbols.c_x => c_x)
    push!(gates, DefaultGates.GateSymbols.c_y => c_y)
    push!(gates, DefaultGates.GateSymbols.c_z => c_z)
    push!(gates, DefaultGates.GateSymbols.r_x => r_x)
    push!(gates, DefaultGates.GateSymbols.r_y => r_y)
    push!(gates, DefaultGates.GateSymbols.r_z => r_z)
    push!(gates, DefaultGates.GateSymbols.r_phase => r_phase)

    push!(gates, DefaultGates.GateSymbols.c_r_x => c_r_x)
    push!(gates, DefaultGates.GateSymbols.c_r_y => c_r_y)
    push!(gates, DefaultGates.GateSymbols.c_r_z => c_r_z)
    push!(gates, DefaultGates.GateSymbols.c_r_phase => c_r_phase)

    push!(gates, DefaultGates.GateSymbols.s => s )
    push!(gates, DefaultGates.GateSymbols.t => t )
    return ;
end

# =========================================================================== #
#                       Custom gate storage: 1 qubit
# =========================================================================== #

"""
    create_gate_1q(gate_label::String, gen_func::Function)

Creates a new user-defined gate, caches the generating function, and returns the
GateOps.GateSymbol key for use in a circuit.

# Examples
```julia-repl

```
"""
function create_gate_1q(gate_label::String, gen_func::Function)
    mat2x2 = nothing
    if applicable(gen_func, pi/2)
        mat2x2 = gen_func(pi/2)
    else
        mat2x2 = gen_func()
    end
    if size(mat2x2) != (2,2)
        error("Please ensure generated matrix is 2x2")
    end
    gs = GateOps.GateSymbol(Symbol(gate_label))
    cache_gate!(gs, gen_func)
    return gs
end

"""
    create_gate_1q(gate_label::String, mat::SArray{Tuple{2,2},Complex{Float64},2,4})

Creates a new user-defined gate, caches the generating matrix, and returns the
GateOps.GateSymbol key for use in a circuit.

# Examples
```julia-repl

```
"""
function create_gate_1q(gate_label::String, mat::SArray{Tuple{2,2},Complex{Float64},2,4})
    return create_gate_1q(gate_label, ()->mat)
end

"""
    create_gate_1q(gate_label::String, mat::Array{Complex{Float64},2})

Creates a new user-defined gate, caches the generating matrix, and returns the
GateOps.GateSymbol key for use in a circuit. Matrix is converted to 2x2 StaticArrays
SMatrix internally.

# Examples
```julia-repl

```
"""
function create_gate_1q(gate_label::String, mat::Array{<:Number,2})
    return create_gate_1q(gate_label, SArray{Tuple{2,2},Complex{Float64},2,4}(mat))
end

# =========================================================================== #
#                       Custom gate storage: 2 qubit
# =========================================================================== #

"""
    create_gate_2q(gate_label::String, gen_func::Function)

Creates a new user-defined gate, caches the generating function, and returns the
GateOps.GateSymbol key for use in a circuit.

# Examples
```julia-repl

```
"""
function create_gate_2q(gate_label::String, gen_func::Function)
    mat4x4 = nothing
    if applicable(gen_func, pi/2)
        mat4x4 = gen_func(pi/2)
    else
        mat4x4 = gen_func()
    end
    if size(mat4x4) != (4,4)
        error("Please ensure generated matrix is 4x4")
    end
    gs = GateOps.GateSymbol(Symbol(gate_label))
    cache_gate!(gs, gen_func)
    return gs
end

"""
    create_gate_2q(gate_label::String, mat::SArray{Tuple{2,2},Complex{Float64},2,4})

Creates a new user-defined gate, caches the generating matrix, and returns the
GateOps.GateSymbol key for use in a circuit.

# Examples
```julia-repl

```
"""
function create_gate_2q(gate_label::String, mat::SArray{Tuple{4,4},Complex{Float64},2,16})
    return create_gate_2q(gate_label, ()->mat)
end
"""
    create_gate_2q(gate_label::String, mat::Array{Complex{Float64},2})

Creates a new user-defined gate, caches the generating matrix, and returns the
GateOps.GateSymbol key for use in a circuit. Matrix is converted to 2x2 StaticArrays
SMatrix internally.

# Examples
```julia-repl

```
"""
function create_gate_2q(gate_label::String, mat::Array{<:Number,2})
    return create_gate_2q(gate_label, SArray{Tuple{4,4},Complex{Float64},2,16}(mat))
end

# =========================================================================== #

"""
    cache_gate!(label::GateOps.AGateSymbol, mat::Matrix{<:Number})

Adds a mapping between label=>mat for fast retrieval of gates in circuit generation.

# Examples
```julia-repl
julia> QXZoo.GateMap.cache_gate!(QXZoo.GateOps.GateSymbol(:mygate), ()->[ 1 0; 0 1])
```
"""
function cache_gate!(key::Union{GateOps.GateSymbol, GateOps.GateSymbolP, Function}, mat_func::Function)
    if ~haskey(gates, key)
        gates[key] = mat_func
    end
    return ;
end

"""
    replace_gate!(label::GateOps.AGateSymbol, mat::Matrix{<:Number})

Replaces existing gate mapping as defined with cache_gate! otherwise performs no-op.

# Examples
```julia-repl
julia> QXZoo.GateMap.replace_gate!(QXZoo.GateOps.GateSymbol(:mygate), ()->[ 1 0; 0 -1])
```
"""
function replace_gate!(key::Union{GateOps.GateSymbol, GateOps.GateSymbolP, Function}, mat_func::Function)
    if haskey(gates, key)
        gates[key] = mat_func
    end
    return ;
end

# =========================================================================== #
#                       Projection operators
# =========================================================================== #
"""
    p00()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the |0><0| projector gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.p00()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 1.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+0.0im
```
"""
function p00()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return @SMatrix [1.0+0.0im 0.0; 0.0 0.0]
end

"""
    p01()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the |0><1| projector gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.p01()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 0.0+0.0im  1.0+0.0im
 0.0+0.0im  0.0+0.0im
```
"""
function p01()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return @SMatrix [0.0 1.0+0.0im; 0.0 0.0]
end

"""
    p10()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the |1><0| projector gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.p10()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 0.0+0.0im  0.0+0.0im
 1.0+0.0im  0.0+0.0im
```
"""
function p10()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return @SMatrix [0.0 0.0; 1.0+0.0im 0.0]
end

"""
    p11()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the |1><1| projector gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.p11()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 0.0+0.0im  0.0+0.0im
 0.0+0.0im  1.0+0.0im
```
"""
function p11()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return @SMatrix [0.0 0.0; 0.0 1.0+0.0im]
end

# =========================================================================== #
#                       Pauli operators
# =========================================================================== #
"""
    x()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the Pauli-X gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.x()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 0.0+0.0im  1.0+0.0im
 1.0+0.0im  0.0+0.0im
```
"""
function x()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return @SMatrix [0.0 1.0+0.0im; 1.0+0.0im 0.0]
end

"""
    y()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the Pauli-Y gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.y()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 0.0+0.0im  -0.0-1.0im
 0.0+1.0im   0.0+0.0im
```
"""
function y()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return @SMatrix [0.0 -1.0im; 1.0im 0.0]
end

"""
    z()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the Pauli-Z gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.z()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 1.0+0.0im   0.0+0.0im
 0.0+0.0im  -1.0+0.0im
```
"""
function z()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return @SMatrix [1.0+0.0im 0.0; 0.0 -1.0+0.0im]
end

# =========================================================================== #
#                       Additional operators
# =========================================================================== #

"""
    h()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the Hadamard gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.h()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 0.707107+0.0im   0.707107+0.0im
 0.707107+0.0im  -0.707107+0.0im
```
"""
function h()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return SMatrix{2,2}([1.0+0.0im 1.0+0.0im; 1.0+0.0im -1.0+0.0im].*(1/sqrt(2)))
end

"""
    I()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the 2x2 identity gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.I()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 1.0+0.0im  0.0+0.0im
 0.0+0.0im  1.0+0.0im
```
"""
function I()::Diagonal{Complex{Float64},SArray{Tuple{2},Complex{Float64},1,2}}
    return StaticArrays.SDiagonal(1.0+0im,1.0+0im)
end

"""
    s()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the S gate (sqrt(Z)).

# Examples
```julia-repl
julia> QXZoo.GateMap.s()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 1.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+1.0im
```
"""
function s()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return sqrt(z())
end

"""
    t()::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the T gate (sqrt(S)).

# Examples
```julia-repl
julia> QXZoo.GateMap.t()
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 1.0+0.0im       0.0+0.0im
 0.0+0.0im  0.707107+0.707107im
```
"""
function t()::SArray{Tuple{2,2},Complex{Float64},2,4}
    return sqrt(s())
end

# =========================================================================== #
#                       Rotation operators
# =========================================================================== #
"""
    r_x(θ::Number)::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the parametric r_x gate: exp(-0.5im*θ*x())

# Examples
```julia-repl
julia> QXZoo.GateMap.r_x(pi/3)
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 0.866025+0.0im       0.0-0.5im
      0.0-0.5im  0.866025+0.0im
```
"""
function r_x(θ::Number)::SArray{Tuple{2,2},Complex{Float64},2,4}
    return exp(-1im*x()*θ/2)
end
"""
    r_y(θ::Number)::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the parametric r_y gate: exp(-0.5im*θ*y())

# Examples
```julia-repl
julia> QXZoo.GateMap.r_y(pi/6)
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 0.965926+0.0im  -0.258819+0.0im
 0.258819+0.0im   0.965926+0.0im
```
"""
function r_y(θ::Number)::SArray{Tuple{2,2},Complex{Float64},2,4}
    return exp(-1im*y()*θ/2)
end
"""
    r_z(θ::Number)::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the parametric r_z gate: exp(-0.5im*θ*z())

# Examples
```julia-repl
julia> QXZoo.GateMap.r_z(pi/2)
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 0.707107-0.707107im       0.0+0.0im
      0.0+0.0im       0.707107+0.707107im
```
"""
function r_z(θ::Number)::SArray{Tuple{2,2},Complex{Float64},2,4}
    return exp(-1im*z()*θ/2)
end
"""
    r_phase(θ::Number)::SArray{Tuple{2,2},Complex{Float64},2,4}

Function which generates the parametric phase shifting gate [1 0; 0 exp(1im*θ)].
Note, this is not the Phase gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.r_phase(pi/4)
2×2 StaticArrays.SArray{Tuple{2,2},Complex{Float64},2,4} with indices SOneTo(2)×SOneTo(2):
 1.0+0.0im       0.0+0.0im
 0.0+0.0im  0.707107+0.707107im
```
"""
function r_phase(θ::Number)::SArray{Tuple{2,2},Complex{Float64},2,4}
    return [1 0; 0 exp(1im*θ)]
end

# =========================================================================== #
#                       Controlled Pauli operators
# =========================================================================== #
"""
    c_x()::SArray{Tuple{4,4},Complex{Float64},2,16}

Function which generates the controlled Pauli-X gate as a block diagonal 4x4 matrix.

# Examples
```julia-repl
julia> QXZoo.GateMap.c_x()
4×4 StaticArrays.SArray{Tuple{4,4},Complex{Float64},2,16} with indices SOneTo(4)×SOneTo(4):
 1.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  1.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  1.0+0.0im
 0.0+0.0im  0.0+0.0im  1.0+0.0im  0.0+0.0im
```
"""
function c_x()::SArray{Tuple{4,4},Complex{Float64},2,16}
    return kron(p00(), I()) + kron(p11(), x())
end

"""
    c_y()::SArray{Tuple{4,4},Complex{Float64},2,16}

Function which generates the controlled Pauli-X gate as a block diagonal 4x4 matrix.

# Examples
```julia-repl
julia> QXZoo.GateMap.c_y()
4×4 StaticArrays.SArray{Tuple{4,4},Complex{Float64},2,16} with indices SOneTo(4)×SOneTo(4):
 1.0+0.0im  0.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  1.0+0.0im  0.0+0.0im  0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.0-1.0im
 0.0+0.0im  0.0+0.0im  0.0+1.0im  0.0+0.0im
```
"""
function c_y()::SArray{Tuple{4,4},Complex{Float64},2,16}
    return kron(p00(), I()) + kron(p11(), y())
end

"""
    c_z()::SArray{Tuple{4,4},Complex{Float64},2,16}

Function which generates the controlled Pauli-X gate as a block diagonal 4x4 matrix.

# Examples
```julia-repl
julia> QXZoo.GateMap.c_z()
4×4 StaticArrays.SArray{Tuple{4,4},Complex{Float64},2,16} with indices SOneTo(4)×SOneTo(4):
 1.0+0.0im  0.0+0.0im  0.0+0.0im   0.0+0.0im
 0.0+0.0im  1.0+0.0im  0.0+0.0im   0.0+0.0im
 0.0+0.0im  0.0+0.0im  1.0+0.0im   0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  -1.0+0.0im
```
"""
function c_z()::SArray{Tuple{4,4},Complex{Float64},2,16}
    return kron(p00(), I()) + kron(p11(), z())
end

# =========================================================================== #
#                       Controlled rotation operators
# =========================================================================== #
"""
    c_r_x(θ::Number)::SArray{Tuple{4,4},Complex{Float64},2,16}

Function which generates the controlled parametric r_x gate as a block diagonal 4x4 matrix.

# Examples
```julia-repl
julia> QXZoo.GateMap.c_r_x(pi/3)
4×4 StaticArrays.SArray{Tuple{4,4},Complex{Float64},2,16} with indices SOneTo(4)×SOneTo(4):
 1.0+0.0im  0.0+0.0im       0.0+0.0im       0.0+0.0im
 0.0+0.0im  1.0+0.0im       0.0+0.0im       0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.866025+0.0im       0.0-0.5im
 0.0+0.0im  0.0+0.0im       0.0-0.5im  0.866025+0.0im
```
"""
function c_r_x(θ::Number)::SArray{Tuple{4,4},Complex{Float64},2,16}
    return kron(p00(), I()) + kron(p11(), r_x(θ))
end

"""
    c_r_y(θ::Number)::SArray{Tuple{4,4},Complex{Float64},2,16}

Function which generates the controlled parametric r_y gate as a block diagonal 4x4 matrix.

# Examples
```julia-repl
julia> QXZoo.GateMap.c_r_y(pi/4)
4×4 StaticArrays.SArray{Tuple{4,4},Complex{Float64},2,16} with indices SOneTo(4)×SOneTo(4):
 1.0+0.0im  0.0+0.0im       0.0+0.0im        0.0+0.0im
 0.0+0.0im  1.0+0.0im       0.0+0.0im        0.0+0.0im
 0.0+0.0im  0.0+0.0im   0.92388+0.0im  -0.382683+0.0im
 0.0+0.0im  0.0+0.0im  0.382683+0.0im    0.92388+0.0im
```
"""
function c_r_y(θ::Number)::SArray{Tuple{4,4},Complex{Float64},2,16}
    return kron(p00(), I()) + kron(p11(), r_y(θ))
end
"""
    c_r_z(θ::Number)::SArray{Tuple{4,4},Complex{Float64},2,16}

Function which generates the controlled parametric r_y gate as a block diagonal 4x4 matrix.

# Examples
```julia-repl
julia> QXZoo.GateMap.c_r_z(pi/6)
4×4 StaticArrays.SArray{Tuple{4,4},Complex{Float64},2,16} with indices SOneTo(4)×SOneTo(4):
 1.0+0.0im  0.0+0.0im       0.0+0.0im            0.0+0.0im
 0.0+0.0im  1.0+0.0im       0.0+0.0im            0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.965926-0.258819im       0.0+0.0im
 0.0+0.0im  0.0+0.0im       0.0+0.0im       0.965926+0.258819im
```
"""
function c_r_z(θ::Number)::SArray{Tuple{4,4},Complex{Float64},2,16}
    return kron(p00(), I()) + kron(p11(), r_z(θ))
end

"""
    c_r_phase(θ::Number)::SArray{Tuple{4,4},Complex{Float64},2,16}

Function which generates the controlled parametric phase shifting gate [1 0; 0 exp(1im*θ)] as a block diagonal 4x4 matrix.

Note, this is not the controlled Phase gate.

# Examples
```julia-repl
julia> QXZoo.GateMap.c_r_phase(pi/9)
4×4 StaticArrays.SArray{Tuple{4,4},Complex{Float64},2,16} with indices SOneTo(4)×SOneTo(4):
 1.0+0.0im  0.0+0.0im  0.0+0.0im       0.0+0.0im
 0.0+0.0im  1.0+0.0im  0.0+0.0im       0.0+0.0im
 0.0+0.0im  0.0+0.0im  1.0+0.0im       0.0+0.0im
 0.0+0.0im  0.0+0.0im  0.0+0.0im  0.939693+0.34202im
```
"""
function c_r_phase(θ::Number)::SArray{Tuple{4,4},Complex{Float64},2,16}
    return kron(p00(), I()) + kron(p11(), r_phase(θ))
end

# =========================================================================== #
"""
Allow access directly using a gate-call for ease-of-use
"""
function Base.getindex(gc::Dict{GateOps.AGateSymbol, Matrix{<:Number}}, gs::GateOps.AGateCall)
    return gc[gs.gate_label]
end

"""
    clear_cache()

Empty the stored gates from the cache.
"""
function clear_cache()
    empty!(gates)
end

end


function LineSearch(arg0::MultivariateOptimizer, arg1::jdouble, arg2::jdouble, arg3::jdouble)
    return LineSearch((MultivariateOptimizer, jdouble, jdouble, jdouble), arg0, arg1, arg2, arg3)
end

function search(obj::LineSearch, arg0::Vector{jdouble}, arg1::Vector{jdouble})
    return jcall(obj, "search", UnivariatePointValuePair, (Vector{jdouble}, Vector{jdouble}), arg0, arg1)
end


function plot(h::Hamiltonian{<:Lattice}; resolution = (1000, 1000), kw...)
    scene = hamiltonianplot(h; resolution = resolution, scale_plot = false, kw...)
    plot = scene[end]
    plot[:tooltips][] && addtooltips!(scene, h)
    scale!(scene)
    return scene
end

@recipe(HamiltonianPlot, hamiltonian) do scene
    Theme(
        allintra = false, allcells = true, showsites = true, showlinks = true,
        shadedsites = false, shadedlinks = false, dimming = 0.75,
        siteradius = 0.2, siteborder = 3, siteborderdarken = 1.0, linkdarken = 0.0,
        linkthickness = 6, linkoffset = missing, linkradius = 0.1,
        tooltips = true, digits = 3,
        _tooltips_rowcolhar = Vector{Tuple{Int,Int,Int}}[],
        colors = map(t -> RGBAf0(t...),
            ((0.960,0.600,.327), (0.410,0.067,0.031),(0.940,0.780,0.000),
            (0.640,0.760,0.900),(0.310,0.370,0.650),(0.600,0.550,0.810),
            (0.150,0.051,0.100),(0.870,0.530,0.640),(0.720,0.130,0.250))),
        light = Vec{3,Float32}[[0, 0, 10], [0, 10, 0], [10, 0, 0], [10, 10, 10], [-10, -10, -10]]
    )
end

function plot!(plot::HamiltonianPlot)
    h = to_value(plot[1])
    lat = h.lattice
    colors = cycle(plot[:colors][])
    sublats = Elsa.sublats(lat)

    mdist = meandist(h)
    mdist > 0 || (mdist = 1)
    plot[:siteradius][] *= mdist/2
    plot[:linkradius][] *= mdist/2

    lo = plot[:linkoffset][]
    plot[:linkoffset][] =
        lo === missing ? √max(1-(plot[:linkradius][]/plot[:siteradius][])^2, 0) : lo * mdist

    # plot sites
    plot[:showsites][] &&
        for (n, har) in enumerate(h.harmonics), (ssrc, csrc) in zip(sublats, colors)
            iszero(har.dn) || plot[:allcells][] || break
            csrc´ = iszero(har.dn) ? csrc : transparent(csrc, 1 - plot[:dimming][])
            itr = siterange(lat, ssrc)
            plotsites!(plot, lat, itr, har.dn, n, csrc´)
        end

    # plot links
    plot[:showlinks][] &&
        for (n, har) in enumerate(h.harmonics)
            iszero(har.dn) || plot[:allcells][] || break
            for (ssrc, csrc) in zip(sublats, colors)
                csrc´ = iszero(har.dn) ? csrc : transparent(csrc, 1 - plot[:dimming][])
                csrc´ = darken(csrc´, plot[:linkdarken][])
                for (sdst, cdst) in zip(sublats, colors)
                    itr = Elsa.eachindex_nz(har, siterange(lat, sdst), siterange(lat, ssrc))
                    plotlinks!(plot, lat, itr, har.dn, n, csrc´)
                end
            end
        end

    return plot
end

function plotsites!(plot, lat, srange, dn, n, color)
    allsites = Elsa.sites(lat)
    br = lat.bravais.matrix
    sites = [padright(allsites[i] + br * dn, Val(3)) for i in srange]
    plot[:tooltips][] && (tt = [(site, 0, n) for site in srange])
    if !isempty(sites)
        plot[:shadedsites][] ? plotsites_hi!(plot, sites, color) : plotsites_lo!(plot, sites, color)
        plot[:tooltips][] && push!(plot[:_tooltips_rowcolhar][], tt)
    end
    return plot
end

function plotsites_lo!(plot, sites, color)
    scatter!(plot, sites;
        color = color,
        markersize = 2 * plot[:siteradius][],
        strokewidth = plot[:siteborder][],
        strokecolor = darken(color, plot[:siteborderdarken][]))
    return nothing
end

function plotsites_hi!(plot, sites, color)
    meshscatter!(plot, sites;
        color = color,
        markersize = plot[:siteradius][], light = plot[:light][])
    return nothing
end

function plotlinks!(plot, lat, itr, dn, n, color)
    links = Pair{SVector{3,Float32},SVector{3,Float32}}[]
    plot[:tooltips][] && (tt = Tuple{Int,Int,Int}[])
    sites = Elsa.sites(lat)
    br = lat.bravais.matrix
    for (row, col) in itr
        iszero(dn) && row == col && continue
        rdst = padright(sites[row] + br * dn, Val(3))
        rsrc = padright(sites[col], Val(3))
        rdst = iszero(dn) ? (rdst + rsrc) / 2 : rdst
        rsrc = rsrc + plot[:siteradius][] * plot[:linkoffset][] * normalize(rdst - rsrc)
        push!(links, rsrc => rdst)
        plot[:tooltips][] && push!(tt, (row, col, n))
    end
    if !isempty(links)
        plot[:shadedlinks][] ? plotlinks_hi!(plot, links, color) : plotlinks_lo!(plot, links, color)
        plot[:tooltips][] && push!(plot[:_tooltips_rowcolhar][], tt)
    end
    return plot
end

function plotlinks_lo!(plot, links, color)
    # linesegments!(plot, links; color = darken(color, plot[:siteborderdarken][]), linewidth = 2+plot[:linkthickness][])
    linesegments!(plot, links; color = color, linewidth = plot[:linkthickness][])
    return nothing
end

function plotlinks_hi!(plot, links, color)
    positions = [(r1 + r2) / 2 for (r1, r2) in links]
    rotvectors = [r2 - r1 for (r1, r2) in links]
    radius = plot[:linkradius][]
    scales = [Vec3f0(radius, radius, norm(r2 - r1)/2) for (r1, r2) in links]
    cylinder = GeometryTypes.Cylinder(Point3f0(0., 0., -1.0), Point3f0(0., 0, 1.0), Float32(1))
    meshscatter!(plot, positions;
        color = color, marker = cylinder, markersize = scales, rotations = rotvectors,
        light = plot[:light][])
    return nothing
end

function addtooltips!(scene, h)
    sceneplot = scene[end]
    visible = Node(false)
    N = Elsa.blockdim(h)
    poprect = lift(scene.events.mouseposition) do mp
        FRect((mp .+ 5), 1,1)
    end
    textpos = lift(scene.events.mouseposition) do mp
        Vec3f0((mp .+ 5 .+ (0, -50))..., 0)
    end
    popup = poly!(campixel(scene), poprect, raw = true, color = RGBAf0(1,1,1,0), visible = visible)
    rect = popup[end]
    translate!(popup, Vec3f0(0, 0, 10000))
    text!(popup, " ", textsize = 30, position = textpos, align = (:center, :center),
        color = :black, strokewidth = 4, strokecolor = :white, raw = true, visible = visible)
    text_field = popup[end]
    on(scene.events.mouseposition) do event
        subplot, idx = mouse_selection(scene)
        layer = findfirst(isequal(subplot), sceneplot.plots)
        if layer !== nothing && idx > 0
            idx´ = fix_linesegments_bug(idx, subplot)
            txt = popuptext(sceneplot, layer, idx´, h)
            text_field[1] = txt
            visible[] = true
        else
            visible[] = false
        end
        return
    end
    return scene
end

# idx returned by mouse_selection seems wrong by a factor 2 in LineSegments subplot
fix_linesegments_bug(idx, subplot::LineSegments) = Int(idx/2) # catches odd idx
fix_linesegments_bug(idx, subplot) = idx

function popuptext(sceneplot, layer, idx, h)
    cache = sceneplot[:_tooltips_rowcolhar][]
    checkbounds(Bool, cache, layer) || return string("Bug in layer: ", layer, " of ", length(cache))
    checkbounds(Bool, cache[layer], idx) || return string("Bug in idx: ", idx, " of ", length(cache[layer]))
    (row, col_or_zero, haridx) = cache[layer][idx]
    if col_or_zero == 0
        col = iszero(col_or_zero) ? row : col_or_zero
        har = h.harmonics[1]
    else
        col = col_or_zero
        har = h.harmonics[haridx]
    end
    element = round.(har.h[row, col], digits = sceneplot[:digits][])
    isreal = all(o -> imag(o) ≈ 0, element)
    txt = isreal ? matrixstring(real.(element)) : matrixstring(element)
    if col_or_zero == 0
        txt´ = string("Onsite[$col] : ", txt)
    else
        txt´ = string("Hopping[$row, $col] : ", txt)
    end
    return txt´
end

matrixstring(x::Number) = string(x)
function matrixstring(s::SMatrix)
    ss = repr("text/plain", s)
    pos = findfirst(isequal('\n'), ss)
    return pos === nothing ? ss : ss[pos:end]
end


"""

```
mddf_frame!(iframe::Int, framedata::FrameData, options::Options, RNG, R::Result)
```

Computes the MDDF for a single frame, modifies the data in the `R` (type `Result`) structure.

"""
function mddf_frame!(iframe::Int, framedata::FrameData, options::Options, RNG, R::Result)

    # Simplify code by assigning some shortened names
    trajectory = framedata.trajectory
    volume_frame = framedata.volume_frame
    rdf_count_random_frame = framedata.rdf_count_random_frame
    md_count_random_frame = framedata.md_count_random_frame
    dc = framedata.dc
    dmin_mol = framedata.dmin_mol
    dref_mol = framedata.dref_mol
    x_solvent_random = framedata.x_solvent_random
    lc_solvent = framedata.lc_solvent
    solute = trajectory.solute
    solvent = trajectory.solvent
    x_solute = trajectory.x_solute
    x_solvent = trajectory.x_solvent

    # Reset counters for this frame
    reset!(volume_frame)
    @. rdf_count_random_frame = 0.0
    @. md_count_random_frame = 0.0

    # get pbc sides in this frame
    sides = getsides(trajectory, iframe)

    volume_frame.total = sides[1] * sides[2] * sides[3]
    R.volume.total = R.volume.total + volume_frame.total

    R.density.solute = R.density.solute + (solute.nmols / volume_frame.total)
    R.density.solvent = R.density.solvent + (solvent.nmols / volume_frame.total)

    # Add the box side information to the box structure, in this frame
    box = Box(options.lcell, sides, R.cutoff)

    # Will wrap everthing relative to the reference atom of the first molecule
    # and move everything such that that center is in the origin. This is important
    # to simplify the computation of cell indexes, as the minimum coordinates are
    # automatically -side/2 at each direction. The reason for wrapping to the solute
    # center and not to the origin is that if the solute is a protein, for example,
    # the protein gets centered, and hardly ever it will be necessary to wrap the
    # coordinates afterwards, accelerating the computation.
    solute_center = x_solute[R.irefatom]
    wrap!(x_solute, sides, solute_center)
    center_to_origin!(x_solute, solute_center)
    wrap!(x_solvent, sides, solute_center)
    center_to_origin!(x_solvent, solute_center)

    # Initialize linked cells
    initcells!(x_solvent, box, lc_solvent)

    # Check if the cutoff is not too large considering the periodic cell size
    if R.cutoff > sides[1] / 2.0 || R.cutoff > sides[2] / 2.0 || R.cutoff > sides[3] / 2.0
        error("""
        in MDDF: cutoff or dbulk > periodic_dimension/2 in frame: $iframe
                 max(cutoff,dbulk) = $(R.cutoff)
                 sides read from file = $(sides)
                 If sides are zero it is likely that the PBC information is not available
                 in the trajectory file.
        """)
    end

    local n_dmin_in_bulk
    n_solvent_in_bulk = 0.0
    for isolute = 1:solute.nmols

        # We need to do this one solute molecule at a time to avoid exploding the memory
        # requirements
        x_this_solute = viewmol(isolute, x_solute, solute)

        # Compute all distances between solute and solvent atoms which are smaller than the
        # cutoff (this is the most computationally expensive part), the distances are returned
        # in the dc structure
        cutoffdistances!(R.cutoff, x_this_solute, x_solvent, lc_solvent, box, dc)

        # For each solute molecule, update the counters (this is highly suboptimal, because
        # within updatecounters there are loops over solvent molecules, in such a way that
        # this will loop with cost nsolute*nsolvent. However, I cannot see an easy solution
        # at this point with acceptable memory requirements
        n_dmin_in_bulk, n_dref_in_bulk =
            updatecounters!(R, solute, solvent, dc, dmin_mol, dref_mol)
        n_solvent_in_bulk += n_dref_in_bulk
    end
    n_solvent_in_bulk = n_solvent_in_bulk / solute.nmols

    #
    # Computing the random-solvent distribution to compute the random minimum-distance count
    #
    bulk_range = (solvent.nmols-n_dmin_in_bulk+1):solvent.nmols
    for i = 1:options.n_random_samples

        # generate random solvent box, and store it in x_solvent_random
        for j = 1:solvent.nmols
            # Choose randomly one molecule from the bulk, if there are actually bulk molecules
            if n_dmin_in_bulk != 0
                jmol = dmin_mol[random(RNG, bulk_range)].jmol
            else
                jmol = random(RNG, 1:solvent.nmols)
            end
            # Indexes of this molecule in the x_solvent array
            x_ref = viewmol(jmol, x_solvent, solvent)
            # Indexes of the random molecule in random array
            x_rnd = viewmol(j, x_solvent_random, solvent)
            # Generate new random coordinates (translation and rotation) for this molecule
            random_move!(x_ref, R.irefatom, sides, x_rnd, RNG)
        end

        # wrap random solvent coordinates to box, with the center at the origin
        wrap!(x_solvent_random, sides)

        # Initialize linked cells
        initcells!(x_solvent_random, box, lc_solvent)

        # Choose randomly one solute molecule to be the solute in this sample
        i_rand_mol = random(RNG, 1:solute.nmols)
        x_this_solute = viewmol(i_rand_mol, x_solute, solute)

        # Compute all distances between solute and solvent atoms which are smaller than the
        # cutoff (this is the most computationally expensive part), the distances are returned
        # in the dc structure
        cutoffdistances!(R.cutoff, x_this_solute, x_solvent_random, lc_solvent, box, dc)

        # Update the counters of the random distribution
        updatecounters!(
            R,
            rdf_count_random_frame,
            md_count_random_frame,
            solvent,
            dc,
            dmin_mol,
            dref_mol,
        )

    end # random solvent sampling

    # Update counters with the data of this frame
    update_counters_frame!(
        R,
        rdf_count_random_frame,
        md_count_random_frame,
        volume_frame,
        solute,
        solvent,
        n_solvent_in_bulk,
    )

    nothing
end


# Copyright (c) 2017: Miles Lubin and contributors
# Copyright (c) 2017: Google Inc.
#
# Use of this source code is governed by an MIT-style license that can be found
# in the LICENSE.md file or at https://opensource.org/licenses/MIT.

"""
    abstract type SetMapBridge{T,S2,S1,F,G} <: AbstractBridge end

Consider two type of sets, `S1` and `S2`, and a linear mapping `A` such that
the image of a set of type `S1` under `A` is a set of type `S2`.

A `SetMapBridge{T,S2,S1,F,G}` is a bridge that maps `G`-in-`S2` constraints
into `F`-in-`S1` by mapping the function through `A`.

The linear map `A` is described by;

 * [`MathOptInterface.Bridges.map_set`](@ref)
 * [`MathOptInterface.Bridges.map_function`](@ref).

Implementing a method for these two functions is sufficient to bridge
constraints. However, in order for the getters and setters of attributes such as
dual solutions and starting values to work as well, a method for the following
functions must be implemented:

 * [`MathOptInterface.Bridges.inverse_map_set`](@ref)
 * [`MathOptInterface.Bridges.inverse_map_function`](@ref)
 * [`MathOptInterface.Bridges.adjoint_map_function`](@ref)
 * [`MathOptInterface.Bridges.inverse_adjoint_map_function`](@ref)

See the docstrings of each function to see which feature would be missing if it
was not implemented for a given bridge.
"""
abstract type SetMapBridge{T,S2,S1,F,G} <: AbstractBridge end

function bridge_constraint(
    BT::Type{<:SetMapBridge{T,S2,S1,F,G}},
    model::MOI.ModelLike,
    func::G,
    set::S1,
) where {T,S2,S1,F,G}
    mapped_func = MOI.Bridges.map_function(BT, func)
    mapped_set = MOI.Bridges.map_set(BT, set)
    constraint = MOI.add_constraint(model, mapped_func, mapped_set)
    return BT(constraint)
end

function MOI.supports_constraint(
    ::Type{<:SetMapBridge{T,S2,S1}},
    ::Type{<:MOI.AbstractScalarFunction},
    ::Type{S1},
) where {T,S2,S1<:MOI.AbstractScalarSet}
    return true
end

function MOI.supports_constraint(
    ::Type{<:SetMapBridge{T,S2,S1}},
    ::Type{<:MOI.AbstractVectorFunction},
    ::Type{S1},
) where {T,S2,S1<:MOI.AbstractVectorSet}
    return true
end

function MOI.Bridges.added_constrained_variable_types(::Type{<:SetMapBridge})
    return Tuple{Type}[]
end

function MOI.Bridges.added_constraint_types(
    ::Type{<:SetMapBridge{T,S2,S1,F}},
) where {T,S2,S1,F}
    return Tuple{Type,Type}[(F, S2)]
end

# Attributes, Bridge acting as a model

function MOI.get(
    ::SetMapBridge{T,S2,S1,F},
    ::MOI.NumberOfConstraints{F,S2},
)::Int64 where {T,S2,S1,F}
    return 1
end

function MOI.get(
    bridge::SetMapBridge{T,S2,S1,F},
    ::MOI.ListOfConstraintIndices{F,S2},
) where {T,S2,S1,F}
    return [bridge.constraint]
end

# References

function MOI.delete(model::MOI.ModelLike, bridge::SetMapBridge)
    MOI.delete(model, bridge.constraint)
    return
end

# Attributes, Bridge acting as a constraint

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintFunction,
    bridge::SetMapBridge{T,S2,S1,F,G},
) where {T,S2,S1,F,G}
    mapped_func = MOI.get(model, attr, bridge.constraint)
    func = MOI.Bridges.inverse_map_function(typeof(bridge), mapped_func)
    return MOI.Utilities.convert_approx(G, func)
end

function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintFunction,
    bridge::SetMapBridge{T,S2,S1,F,G},
    func::G,
) where {T,S2,S1,F,G}
    new_f = MOI.Bridges.map_function(typeof(bridge), func)
    MOI.set(model, attr, bridge.constraint, new_f)
    return
end

function MOI.get(
    model::MOI.ModelLike,
    attr::MOI.ConstraintSet,
    bridge::SetMapBridge,
)
    set = MOI.get(model, attr, bridge.constraint)
    return MOI.Bridges.inverse_map_set(typeof(bridge), set)
end

function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintSet,
    bridge::SetMapBridge{T,S2,S1},
    set::S1,
) where {T,S2,S1}
    new_set = MOI.Bridges.map_set(typeof(bridge), set)
    MOI.set(model, attr, bridge.constraint, new_set)
    return
end

function MOI.supports(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintPrimalStart,MOI.ConstraintDualStart},
    ::Type{<:SetMapBridge{T,S2,S1,F}},
) where {T,S2,S1,F}
    return MOI.supports(model, attr, MOI.ConstraintIndex{F,S2})
end

function MOI.get(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintPrimal,MOI.ConstraintPrimalStart},
    bridge::SetMapBridge,
)
    value = MOI.get(model, attr, bridge.constraint)
    return MOI.Bridges.inverse_map_function(typeof(bridge), value)
end

function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintPrimalStart,
    bridge::SetMapBridge,
    value,
)
    mapped_value = MOI.Bridges.map_function(typeof(bridge), value)
    MOI.set(model, attr, bridge.constraint, mapped_value)
    return
end

function MOI.get(
    model::MOI.ModelLike,
    attr::Union{MOI.ConstraintDual,MOI.ConstraintDualStart},
    bridge::SetMapBridge,
)
    value = MOI.get(model, attr, bridge.constraint)
    return MOI.Bridges.adjoint_map_function(typeof(bridge), value)
end

function MOI.set(
    model::MOI.ModelLike,
    attr::MOI.ConstraintDualStart,
    bridge::BT,
    value,
) where {BT<:SetMapBridge}
    mapped_value = MOI.Bridges.inverse_adjoint_map_function(BT, value)
    MOI.set(model, attr, bridge.constraint, mapped_value)
    return
end

# By linearity of the map, we can just change the constant/coefficient

function MOI.modify(
    model::MOI.ModelLike,
    bridge::BT,
    change::MOI.ScalarConstantChange,
) where {BT<:SetMapBridge}
    new_constant = MOI.Bridges.map_function(BT, change.new_constant)
    MOI.modify(model, bridge.constraint, MOI.ScalarConstantChange(new_constant))
    return
end

function MOI.modify(
    model::MOI.ModelLike,
    bridge::BT,
    change::MOI.VectorConstantChange,
) where {BT<:SetMapBridge}
    new_constant = MOI.Bridges.map_function(BT, change.new_constant)
    MOI.modify(model, bridge.constraint, MOI.VectorConstantChange(new_constant))
    return
end

function MOI.modify(
    model::MOI.ModelLike,
    bridge::BT,
    change::MOI.ScalarCoefficientChange,
) where {BT<:SetMapBridge}
    new_coefficient = MOI.Bridges.map_function(BT, change.new_coefficient)
    new_change = MOI.ScalarCoefficientChange(change.variable, new_coefficient)
    MOI.modify(model, bridge.constraint, new_change)
    return
end

function MOI.modify(
    model::MOI.ModelLike,
    bridge::BT,
    change::MOI.MultirowChange,
) where {BT<:SetMapBridge}
    # It is important here that `change.new_coefficients` contains the complete
    # new sparse column associated to the variable. Calling modify twice with
    # part of the column won't work since  the linear map might reset all the
    # column each time.
    coefficients = MOI.Bridges.map_function(BT, change.new_coefficients)
    new_change = MOI.MultirowChange(change.variable, coefficients)
    MOI.modify(model, bridge.constraint, new_change)
    return
end


# This file is a part of BAT.jl, licensed under the MIT License (MIT).

using .NestedSamplers               # used in this file, because this file ist only load if it is required

include("ens_bounds.jl")
include("ens_proposals.jl")

"""
    struct EllipsoidalNestedSampling <: AbstractSamplingAlgorithm

*Experimental feature, not part of stable public API.*

Uses the julia package
[NestedSamplers.jl](https://github.com/TuringLang/NestedSamplers.jl) to use nested sampling algorithm.

Constructors:

* ```$(FUNCTIONNAME)(; fields...)```

Fields:

$(TYPEDFIELDS)

!!! note

    This functionality is only available when the
    [NestedSamplers.jl](https://github.com/TuringLang/NestedSamplers.jl) package
    is loaded (e.g. via
    `import`).
"""
@with_kw struct EllipsoidalNestedSampling{TR<:AbstractDensityTransformTarget} <: AbstractSamplingAlgorithm
    trafo::TR = PriorToUniform()

    "Number of live-points."
    num_live_points::Int = 1000

    "Volume around the live-points."
    bound::ENSBound = ENSEllipsoidBound()

    "Algorithm used to choose new live-points."
    proposal::ENSProposal = ENSAutoProposal()

    "Scale factor for the volume."
    enlarge::Float64 = 1.25

    # "Not sure about how this works yet."
    # update_interval::Float64 =

    "Number of iterations before the first bound will be fit."
    min_ncall::Int64 = 2*num_live_points

    "Efficiency before fitting the first bound."
    min_eff::Float64 = 0.1

    # "The following four are the possible convergence criteria to end the algorithm."
    dlogz::Float64 = 0.01
    max_iters = Inf
    max_ncalls = 10^7
    maxlogl = Inf
end
export EllipsoidalNestedSampling

function bat_sample_impl(rng::AbstractRNG, target::AnyDensityLike, algorithm::EllipsoidalNestedSampling)

    density_notrafo = convert(AbstractDensity, target)
    shaped_density, trafo = bat_transform(algorithm.trafo, density_notrafo)                 # BAT prior transformation
    vs = varshape(shaped_density)
    density = unshaped(shaped_density)
    dims = totalndof(vs)

    model = NestedModel(logdensityof(density), identity);                                   # identity, because ahead the BAT prior transformation is used instead
    bounding = ENSBounding(algorithm.bound)
    prop = ENSprop(algorithm.proposal)
    sampler = Nested(dims, algorithm.num_live_points;
                        bounds=bounding, proposal=prop,
                        enlarge=algorithm.enlarge, min_ncall=algorithm.min_ncall, min_eff=algorithm.min_eff
    )

    samples_w, state = sample(model, sampler;                                               # returns samples with weights as one vector and the actual state
        dlogz = algorithm.dlogz, maxiter = algorithm.max_iters,
        maxcall = algorithm.max_ncalls, maxlogl = algorithm.maxlogl, chain_type=Array
    )

    weights = samples_w[:, end]                                                             # the last elements of the vectors are the weights
    nsamples = size(samples_w,1)
    samples = [samples_w[i, 1:end-1] for i in 1:nsamples]                                   # the other ones (between 1 and end-1) are the samples
    logvals = map(logdensityof(density), samples)                                           # posterior values of the samples
    samples_trafo = vs.(BAT.DensitySampleVector(samples, logvals, weight = weights))
    samples_notrafo = inverse(trafo).(samples_trafo)                                            # Here the samples are retransformed

    logintegral = Measurements.measurement(state.logz, state.logzerr)
    ess = bat_eff_sample_size(samples_notrafo, KishESS()).result

    return (
        result = samples_notrafo, result_trafo = samples_trafo, trafo = trafo,
        logintegral = logintegral, ess = ess,
        info = state
    )
end