text:
"""
    MTPD_variational(mt_enc, mt_post, hphi, emission_coeffs, logstd, d, flag_mt_x0)

A deterministic MT Pharmacodynamic model (subtype of `MTDSModel`) following the equations:

`` [β₁, β₂, β₃, θ₁, ..., θd, α] = h_ϕ(z)``

`` x_t = diag(β₁) x_{t-1} + β₂ u_{t} ``

`` y_t = g_θ(x_t + β₃) + α + s .* ϵ_t``

for ``ϵ_t \\sim 𝑵(0, I)`` and ``x_0 = 0``. This is a discrete time version of the usual
pharmacodynamic model using an effect site per dimension, but with a more general sigmoidal emission
function than simply the Hill Equation (and also more numerically stable to boot). Notable
characteristics:

1. Latent system is deterministic.
2. The latent state ``x_0`` is constant and equal to zero.
3. The emission covariance is diagonal.

## Struct description
The model is variational in ``z``, and contains an encoder (`mt_enc`) which is
an RNN which encodes the sequences `Y` and `U` (which should be of size ``d × T ×
n_b``, where ``d`` is the size of each vector, ``T`` is the sequence length, and
``n_b`` is the batch size) into a single vector. One can instead perform standard
(non-amortized VI by using a LookupTable here. The `mt_post` network is a
network which returns a tuple of (mean, (diagonal) sqrt variance) defining a
variational Gaussian posterior. The method `encode` will return two tuples, the
first of which is a sample from this posterior (using reparameterization),
elements 2 and 3 are the variational posterior mean and stdev components. (The
second tuple is empty -- this is reserved for ``x_0`` if it has a separate
encoder.) Note that these variational networks need not be used, e.g. if one is
using an alternative MCO approach.

The function `hphi` is a network ``h_ϕ: \\text{dom}(z) → \\text{dom}(θ)`` which
takes a ``k × n_b`` matrix ``z`` and returns the LDS parameter vector θ. The emission
function is a flexible sigmoidal function, a kind of 1-hidden layer MLP with fixed
parameters (``emission\\_coeffs``) for the first layer. ``d`` is both the dimension of ``x``
and ``y``, since there is one latent dimension per output dimension.
``emission\\_coeffs`` is an L × 2 matrix (columns give a, b respectively, see paper),
and finally  `flag_mt_x0` is a placeholder for whether the ``x_0`` is modulated via
``z`` (not yet implemented).

## Methods
As per all `MTDSModel` objects, the following interface is available:

- `encode(m::MTPD_variational, Y::AbstractArray{T}, U::AbstractArray{T})`,
  which encodes the sequences `Y` and `U` into a sample and variational
  posterior parameters of ``z``.
- `forward(m::MTPD_variational, z::AbstractVecOrMat{T}, U::AbstractArray{T})`,
  the forward model, which constructs the LDS for the latent variable `z` and
  calculates the output mean.
- `reconstruct(m::MTDSModel, Y, U)` which performs both `encode` and `forward`
  operations, effectively reconstructing `Y` through the bottleneck of the
  inference of ``z``.

### Constructor
See `create_mtpd` for a more helpful constructor.

"""
struct MTPD_variational{U,V,W,S,T} <: MTDSModel
    mt_enc::U
    mt_post::V
    hphi::W
    emission_coeffs::S
    logstd::T
    d::Int
    flag_mt_x0::Bool
end

Flux.@treelike MTPD_variational
# common superclass properties
is_amortized(m::MTPD_variational) = true
is_amortized(m::MTPD_variational{U,V,W,S,T}) where {U <: LookupTable,V,W,S,T} = false
has_x0_encoding(m::MTPD_variational) = false

getL(m::MTPD_variational) = size(m.emission_coeffs, 1)

function Base.show(io::IO, l::MTPD_variational)
    d_x0, d_mt = l.d, size(l.mt_post.Dense1.W, 1)
    d_uy = format("effect_sites={:d}", l.d)
    h0_type = l.flag_mt_x0 ? "'MT'" : "'non-MT'"
    enc_type = l.mt_enc isa LookupTable ? "'Standard VI'" : "'Amortized VI'"
    print(io, "MTPD_variational(", d_uy, ", d_mt=", d_mt,
    ", x0_type=", h0_type,  ", enc_type=", enc_type,")")
end

function f_theta_unpack(θ::AbstractVector, m)
    dx = dy = m.d
    L = getL(m)
    βs = σ.(θ[1:dx]), softplus.(θ[(dx+1):2dx]), θ[(2dx+1):3dx]
    ηs = [-softplus.(θ[(3dx+(i-1)*L+1):(3dx+i*L)]) for i in 1:dy]
    α = θ[(3dx+L*dy+1):(3dx+L*dy+dy)]
    return βs, ηs, α
end
function f_theta_unpack(θ::AbstractMatrix, m)
    dx = dy = m.d
    L = getL(m)
    βs = σ.(θ[1:dx, :]), softplus.(θ[(dx+1):2dx, :]), θ[(2dx+1):3dx, :]
    ηs = [-softplus.(θ[(3dx+(i-1)*L+1):(3dx+i*L), :]) for i in 1:dy]
    α = θ[(3dx+L*dy+1):(3dx+L*dy+dy), :]
    return βs, ηs, α
end

forward(m::MTPD_variational, z, U; T_steps=size(U, 2)) = forward(m, nothing, z, U; T_steps=T_steps)
function forward(m::MTPD_variational, x0, z::AbstractVecOrMat{T}, U::AbstractArray{T}; T_steps=size(U, 2)) where T <: AbstractFloat
    # note that T_steps is currently ignored (this is no great issue).
    !(x0 === nothing) && @warn "MTPD_variational currently doesn't accept x0 -- fixed at 0." *
        "Ignoring..."
    @argcheck size(z, 2) == size(U, 3)
    θ = m.hphi(z)
    # currently don't modulate the initial state.
    return _mtpd_model_forward_parvector(m, θ, U)
end

    function _mtpd_model_forward_parvector(m::MTPD_variational, θ::AbstractVector{T},
            U::AbstractMatrix{T}, h0::AbstractVector{T}=zeros(T, m.d)) where T <: AbstractFloat
        n_u, tT = size(U)
        βs, ηs, α = f_theta_unpack(θ, m)

        lds = Flux.Recur(LDSCell_diag_u(βs[1], reshape(βs[2], m.d, n_u), h0))
        lds_state = [lds(U[:,t]) for t in 1:tT]
        X = hcat(lds_state...) .+ βs[3]
        a, b = m.emission_coeffs[:,1], m.emission_coeffs[:,2]
        # h = kron(X, a) - repeat(b, m.n_y)
        yhat = map(1:m.d) do j
            unsqueeze(ηs[j],1) * σ.(a .* X[j:j,:] .- b) .+ α[j]
        end
        return vcat(yhat...)
    end

    function _mtpd_model_forward_parvector(m::MTPD_variational, θ::AbstractMatrix{T},
            U::AbstractArray{T,3}, h0::AbstractVecOrMat{T}=zeros(T, m.d)) where T <: AbstractFloat
        n_u, tT, n_b = size(U)
        h0 = ensure_matrix(h0)
        βs, ηs, α = f_theta_unpack(θ, m)

        lds = Flux.Recur(LDSCell_diag_batch_u(βs[1],
                                              reshape(βs[2], m.d, n_u, n_b),
                                              h0))
        lds_state = [unsqueeze(lds(U[:,t,:]),2) for t in 1:tT]
        X = hcat(lds_state...) .+ unsqueeze(βs[3], 2)
        a, b = m.emission_coeffs[:,1], m.emission_coeffs[:,2]

        yhat = map(1:m.d) do j
            batch_matmul(unsqueeze(ηs[j],1), σ.(a .* X[j:j,:,:] .- b)) .+ unsqueeze(α[j:j,:], 2)
        end
        return vcat(yhat...)
    end

encode(m::MTPD_variational, Y::AbstractArray{T}, U::AbstractArray{T}; T_enc=size(Y, 2),
    stochastic=true) where T = encode(m, vcat(Y, U); T_enc=T_enc, stochastic=stochastic)

encode(m::MTPD_variational, Y::MaskedArray{T}, U::AbstractArray{T}; T_enc=size(Y, 2),
    stochastic=true) where T =
    encode(m, vcat(stacked(Y), U); T_enc=T_enc, stochastic=stochastic)

function encode(m::MTPD_variational, data_enc::AbstractArray; T_enc=size(data_enc, 2), stochastic=true)
    m.flag_mt_x0 && error("Unable to modulate x0 currently. Change encoder code, and `_mtlds_model_forward`.")
    smpz, μ_z, σ_z = posterior_sample(m.mt_enc, m.mt_post, data_enc, T_enc, stochastic)
    return (smpz, μ_z, σ_z), (nothing, nothing, nothing)   # second tuple is for if x0 has separate encoder.
end

#  Cannot define function on an abstract type, so we have to def these inside each subtype.
(m::MTPD_variational)(z::AbstractArray{T,2}, U::AbstractArray{T,3}) where T = forward(m, z, U)
(m::MTPD_variational)(Y::AbstractArray{T,3}, U::AbstractArray{T,3}) where T = reconstruct(m, Y, U)

"""
    create_mtpd(d, d_in, d_enc_state, d_mt; encoder=:LSTM,
    d_hidden=64, hphi=:MLP, emission_pars=:default)

Create `MTPD_variational` with ``d`` effect sites, and hence the dimension of both
the latent system and observations is ``d``. ``d_in`` specifies the dimension of
the input vectors, `d_enc_state` the latent dim of the RNN encoder, `d_mt` of the
latent hierarchical variable ``z``.

### Optional arguments:

- `encoder` is the encoding RNN for the variational posterior. Can be in {`:RNN`
, `:GRU`, `:LSTM`, `:LookupTable`}, with the latter providing standard VI.
- `d_hidden`: number of hidden units in the one-hidden-layer MLP used for ``h_ϕ``.
- `spherical_var`: whether the emission covariance is spherical (and hence has a
  single parameter) (default: false), or is an axis-aligned ellipsoid with `d_y`
  parameters.
- `hphi`: default `:MLP`, can be `:Linear`. This determines whether the prior over
  θ is linear or nonlinear. If == `:MLP`, `d_hidden` specifies layers, o.w. it's ignored.
"""
function create_mtpd(d, d_in, d_enc_state, d_mt; encoder=:LSTM, d_hidden=64,
    spherical_var=false, hphi=:MLP, emission_coeffs=:default, indep_biases::Bool=false)

    # Encoding z from the U, Y sequences
    if encoder == :LSTM
        init_enc = LSTM(d + d_in, d_enc_state)
    elseif encoder == :GRU
        init_enc = GRU(d + d_in, d_enc_state)
    elseif encoder == :RNN
        init_enc = RNN(d + d_in, d_enc_state)
    elseif encoder == :LookupTable
        init_enc = LookupTable(d_mt, 0.01f0, -2.0f0)  # 2/3: initial posterior noise for mu/logstd
    else
        error("encoder must be specified as :LSTM, :GRU, :RNN, :LookupTable")
    end

    # Latent z *encoding* --> posterior E(z|Y,U) and V(z|Y,U)
    enc_post = MultiDense(Dense(d_enc_state, d_mt, identity), Dense(d_enc_state, d_mt, σ))
    enc_post.Dense2.b.data .= -2   # initialize posteriors to be low variance

    # Emission function first layer
    nl_ab =  [0.0   0.008  0.0259    0.1248  0.6012   0.1257  0.0261  0.008;
              0.0  -500.0  -150.0  -18.1187  0.4374  18.8683   150.0  500.0]' |> f32 |> Matrix
    emission_coeffs = (emission_coeffs == :default) ? nl_ab : emission_coeffs
    @assert emission_coeffs isa AbstractMatrix
    @assert size(emission_coeffs, 2) == 2 "Expecting an L×2 matrix for emission_coeffs or :default."
    L = size(emission_coeffs, 1)   # number of basis functions for emission fn.

    # Multi-task network z --> θ
    n_dec = indep_biases ? 3*d + d*L : 3*d + d*L + d
    d_mt_h = indep_biases ? d_mt - d : d_mt
    @assert d_mt_h >= 0 "Independent biases selected, but d_mt is too small to permit this."
    # -- Major mapping z --> θ
    if hphi == :MLP
        hphi = mlp(d_mt_h, d_hidden..., Int(n_dec))
    elseif hphi == :Linear
        hphi = Dense(d_mt_h, Int(n_dec))
    else
        error("hphi must be specified as :MLP, :Linear")
    end
    # -- z --> biases if independent
    if indep_biases
        k_not_bias, k_bias = 1:d_mt_h, (d_mt_h+1):d_mt
        hphi = PartialSplit(k_not_bias, k_bias, 1, hphi, identity)
    end

    # Emission noise standard deviation: either same per channel (spherical) or not (diagonal).
    emission_std = spherical_var ? Tracker.param([0.0f0]) : Tracker.param(zeros(Float32, d))

    MTPD_variational(init_enc, enc_post, hphi, emission_coeffs, emission_std, d, false)
end


text:
module ValueHistoriesLH

using DataStructures, DocStringExtensions, JSON3
# using RecipesBase

import Base: isempty

export
    ValueHistory,
      UnivalueHistory,
        History,
        QHistory,
      MultivalueHistory,
        MVHistory,
        increment!,
    @trace

# MVHistory
export has_series, series_names
export history_to_dict, retrieve, load_history, save_history, load_mvhistory, save_mvhistory
# Single valued histories
export indices

include("abstract.jl")
include("history.jl")
include("qhistory.jl")
include("mvhistory.jl")
include("load_save.jl")
# include("recipes.jl")

end # module

ext: jl


text:
# --------------------------------------------------------------------------- #
# General image type

abstract type Image end

# --------------------------------------------------------------------------- #
# Greyscale Images

struct GreyscaleImage <: Image
  data :: Array{Float32, 2}
  GreyscaleImage(data) = new(autoscale(convert(Array{Float32}, data)))
end

function GreyscaleImage(data :: Array{Float32, 3})
  @assert size(data, 3) == 1
  GreyscaleImage(reshape(data, size(data)[1:2]))
end

# --------------------------------------------------------------------------- #
# RGB Images

struct RGBImage <: Image
  data :: Array{Float32, 3}
  RGBImage(data) = new(autoscale(convert(Array{Float32}, data)))
end

function greyscale(img :: RGBImage)
  # Using rec 601 luma standard
  data = 0.299img.data[:,:,1] + 0.587img.data[:,:,2] + 0.114img.data[:,:,3]
  return GreyscaleImage(data)
end

# --------------------------------------------------------------------------- #
# Create datasets from high-res images

function intervals(L, l, m)
  grid = Int[]
  i = 1

  while i + l - 1 <= L
    push!(grid, i)
    i += l + m
  end

  return grid
end

function ordered_patches{I <: Image}(img :: I, lbl;
                                     size = (128, 128),
                                     offset = nothing,
                                     margin :: Integer = 0,
                                     shuffle :: Bool = true,
                                     imageop = Id(),
                                     multitude = 1)

  # Note: components of argument size should be divisible by 8,
  # otherwise some (multiresolution) networks will not like the input

  n, m = imgsize(img)

  @assert (n >= size[1] && m >= size[2])
  @assert (-margin < min(size...))

  ygrid = intervals(n, size[1], margin)
  xgrid = intervals(m, size[2], margin)

  if offset == nothing
    dy = div(n - (ygrid[end] + size[1] - 1), 2)
    dx = div(m - (xgrid[end] + size[2] - 1), 2)
    offset = (dy, dx)
  end

  ygrid = Int[ i + offset[1] for i in ygrid ]
  xgrid = Int[ j + offset[2] for j in xgrid ]

  imgs = I[]
  lbls = Label[]
  ld = convert(Matrix{Int}, lbl)

  for i in ygrid, j in xgrid
    ii = i + size[1] - 1
    jj = j + size[2] - 1
    push!(imgs, I(img.data[i:ii, j:jj, :]))
    if !isempty(lbl)
      sel = (i .<= ld[2,:] .<= ii) .& (j .<= ld[1,:] .<= jj)
      push!(lbls, Label(ld[:,sel] .- [j-1, i-1]))
    else
      push!(lbls, Label())
    end
  end

  data = vcat( (apply.(imageop, imgs, lbls) for i in 1:multitude)... )
  imgs = [ img for (img, _) in data ]
  lbls = [ lbl for (_, lbl) in data ]

  return imgs, lbls
end

function random_patches{I <: Image}(img :: I, lbl, number;
                                    size = (128, 128),
                                    imageop = Id(),
                                    multitude = 1)

  # Note: components of argument size should be divisible by 8,
  # otherwise some (multiresolution) networks will not like the input

  n, m = imgsize(img)
  @assert (n >= size[1] && m >= size[2])

  imgs = I[]
  lbls = Label[]
  ld = convert(Matrix{Int}, lbl)

  for k in 1:number
    i, j = rand(1:n-size[1]), rand(1:m-size[2])
    ii = i + size[1] - 1
    jj = j + size[2] - 1
    push!(imgs, I(img.data[i:ii, j:jj, :]))
    if !isempty(lbl)
      sel = (i .<= ld[2,:] .<= ii) .& (j .<= ld[1,:] .<= jj)
      push!(lbls, Label(ld[:,sel] .- [j-1, i-1]))
    else
      push!(lbls, Label())
    end
  end

  data = vcat( (apply.(imageop, imgs, lbls) for i in 1:multitude)... )
  imgs = [ img for (img, _) in data ]
  lbls = [ lbl for (_, lbl) in data ]

  return imgs, lbls
end

function ordered_patches(img; kwargs...)
  return ordered_patches(img, Label(); kwargs...)[1]
end

function random_patches(img, number; kwargs...)
  return random_patches(img, Label(), number; kwargs...)[1]
end

# --------------------------------------------------------------------------- #
# Convenience functions

imgdata(img :: Image) = img.data
imgsize(a) = size(a)[1:2]
imgtype{I<:Image}(a::I) = I

imgchannels(::GreyscaleImage) = 1
imgchannels(::RGBImage) = 3
imgchannels(::Type{<: GreyscaleImage}) = 1
imgchannels(::Type{<: RGBImage}) = 3

Base.size(img::Image, args...; kwargs...) = size(img.data, args..., kwargs...)

function imgconvert{T <: Images.Gray}(imgarray :: Array{T})
  data = permutedims(Images.channelview(imgarray), (2, 3, 1))
  return GreyscaleImage(reshape(data, size(data)[1:2]...))
end

function imgconvert{T <: Images.RGB}(imgarray :: Array{T})
  data = permutedims(Images.channelview(imgarray), (2, 3, 1))
  return RGBImage(data)
end

function crop(data :: Matrix{Float32}, x :: Int, y :: Int, w :: Int, h :: Int)
  return data[y:y+h-1,x:x+w-1]
end

function crop(img :: RGBImage, x :: Int, y :: Int, w :: Int, h :: Int)
  return RGBImage(imgdata(img)[y:y+h-1,x:x+w-1,:])
end

function crop(img :: GreyscaleImage, x :: Int, y :: Int, w :: Int, h :: Int)
  return GreyscaleImage(imgdata(img)[y:y+h-1,x:x+w-1])
end

Base.similar(img::GreyscaleImage) = GreyscaleImage(zeros(Float32, imsize(img)))

max_stars_repo_name: tscode/DCellC

max_issues_repo_name: tscode/DCellC


text:
module GeneratorsX

# Use README as the docstring of the module:
@doc let path = joinpath(dirname(@__DIR__), "README.md")
    include_dependency(path)
    replace(read(path, String), r"^```julia"m => "```jldoctest README")
end GeneratorsX

export @generator, @yield

import ContextualMacros
using AbstractYieldMacros: @yield, @yieldfrom
using Base.Meta: isexpr
using FGenerators: define_foldl
using IRTools
using IRTools: @dynamo, IR, Statement, argument!, arguments, functional, return!, xcall
using MacroTools: @capture, combinedef, prewalk, splitdef
using Transducers: Transducers

include("utils.jl")
include("core.jl")

end # module


text:
using Dijkstra
using Test

# Create new graph. By default, labels are of type Integer
G = Graph()

# Edges are added to graph using function `add_edge!`. First two parameters
# after G are node labels and third one is the edge weight.
add_edge!(G, 1, 2, 3)
add_edge!(G, 1, 3, 2)

# To get edge weights back, there's a function `get_edge_weight`:
@test get_edge_weight(G, 1, 2) == 3
@test get_edge_weight(G, 1, 3) == 2

# Graph also stores adjacency lists, and adjacent nodes can be returned using
# `get_adjacent_nodes`:
@test get_adjacent_nodes(G, 1) == Set([2, 3])

# Random graphs can be generated using `generate_random_graph`. For example,
# to generate graph with 10 nodes, edge density 0.5 and smallest weight 1 and
# largest weight 10:

using Random
Random.seed!(1)
G = generate_random_graph(10, 0.5, 1, 10)
@test get_adjacent_nodes(G, 1) == Set([7, 4, 10, 2, 3])


text:
 # compute the survivor probability function using standard (slow) method. Useful to compare results with correlation function
 function corr_bench(traj::Trajectory,cutoff::Float64)
   # main function - calculation of the matrix distance
   dist, domain, evals, sp, nprot, nsvt, nframes, stime = autocorr(traj)
   # evaluation if the solvent is inside (or not) of the "protein domain"
   domain_eval!(dist,domain,cutoff)
   # evaluation of the number of events existing in the domain matrix
   events!(domain,evals,nframes,nsvt)
   # sample space
   sample_space!(sp, domain, nsvt, nframes)
   prob_sur = survivor_prob(evals, sp, nsvt)
   return stime, prob_sur
 end

 # Calculation of survivor probability function using CellListMap.jl package
 function correlation(traj::Trajectory,cutoff::Float64)
   # main function - calculation of the matrix distance
   dist, domain, evals, sp, nprot, nsvt, nframes, stime = autocorr_cell(traj)
   # evaluation if the solvent is inside (or not) of the "protein domain"
   domain_eval!(dist,domain,cutoff)
   # evaluation of the number of events existing in the domain matrix
   events!(domain,evals,nframes,nsvt)
   # sample space
   sample_space!(sp, domain, nsvt, nframes)
   prob_sur = survivor_prob(evals, sp, nsvt)
   return stime, prob_sur
 end


text:
# ---------------------------------------------------------------------------- #
#
#   wrapEgads.jl
#
#   Generates the files in wrapped_api/.
#
#   This script follows the C-wrapping of the Sundials.jl package
#
#   This script is not an active part of the package.
#   The Clang.jl package is used to wrap the egads API.
#
#   To run the script from Julia console:
#     include(joinpath(Pkg.dir("egads"), "src", "wrap_egads.jl"));
#   Note that you really should not have to do this, this script is just
#   provided for completeness.
#
#   egads.jl
#   Spring 2018
#
#   Max Opgenoord
#
# ---------------------------------------------------------------------------- #

using Clang.wrap_c

# Input parameters
const egadspath = "/home/mopg/code/EngSketchPad"

# `outpath` specifies, where the julia wrappers would be generated.
# If the generated .jl files are ok, they have to be copied to the "wrapped_api"
# in the "src/" folder overwriting the old ones
const outpath = normpath( joinpath(dirname(@__FILE__), "new") )
mkpath(outpath)

# Find all relevant egads headers
const incpath = normpath( joinpath(egadspath, "include") )
if !isdir(incpath)
    error("Egads C headers not found. You need to install Engineering Sketch Pad (ESP) first")
end

info("Scanning EGADS headers in $incpath...")

const egads_headers = ["common.h","egads.h","egadsBlendSens.h","egadsErrors.h",
                          "egadsSplineFit.h","egadsTris.h",
                          "egadsTypes.h","egg.h","emp.h","prm.h"]

kk = 1
for folder in egads_headers
    egads_headers[kk] = joinpath( incpath, egads_headers[kk] )
    kk += 1
end

const clang_path = "/usr/lib/clang/3.8.0" # change to your clang location
const clang_includes = [
    joinpath(clang_path, "include"),
]

# Callback to test if a header should actually be wrapped (for exclusion)
function wrap_header(top_hdr::AbstractString, cursor_header::AbstractString)
    (top_hdr == cursor_header) # don't wrap if header is included from the other
                               #   header (e.g. nvector in cvode or cvode_direct
                               #   from cvode_band)
end

function julia_file(header::AbstractString)
    src_name = basename(dirname(header))
    if src_name == "egads"
        src_name = "libegads" # avoid having both egads.jl and egads.jl
    end
    return joinpath(outpath, string(src_name, ".jl"))
end
function library_file(header::AbstractString)
    header_name = basename(header)
    return string("libegads_", basename(dirname(header)))
end

const contxt = wrap_c.init(
    common_file = joinpath(outpath, "types_and_consts.jl"),
    clang_diagnostics = true,
    clang_includes = [clang_includes; incpath],
    header_outputfile = julia_file,
    header_library = library_file,
    header_wrapped=wrap_header
)
contxt.headers = egads_headers

run(contxt)

# NOTE: in the generated files
#   $(Expr(:typealias, :ego, :(Ptr{egObject})))
# needs to be replaced with
#   const ego = Ptr{egObject}
# and something is complaining about FILE, just comment it out
#
# Furthermore, the const error codes need to Int32


text:
using Test
using LibSymspg

include("test_api.jl")

@testset "Cell standardize(export)" begin
    latt = [4.0 0.0 0.0;
            0.0 4.0 0.0;
            0.0 0.0 4.0]
    positions = [0. 0. 0.; 0.5 0.5 0.5]
    types = [1, 1]
    new_latt, new_positions, new_types = find_primitive(latt, positions, types, 1e-5)
    # test arguments not modified
    @test latt ≈ [4.0 0.0 0.0; 0.0 4.0 0.0; 0.0 0.0 4.0]
    @test positions ≈ [0. 0. 0.; 0.5 0.5 0.5]
    @test types == [1, 1]

    @test new_latt ≈ [-2.0 2.0 2.0; 2.0 -2.0 2.0; 2.0 2.0 -2.0]
    @show new_positions
    @test new_positions ≈ [0.0 0.0 0.0]
    @test new_types == [1]

    latt = [-2.0 2.0 2.0; 2.0 -2.0 2.0; 2.0 2.0 -2.0]
    positions = Array{Float64, 2}([0.0 0.0 0.0])
    types = [1]
    latt, positions, types = refine_cell(latt, positions, types, 1e-5)
    @test latt ≈ [4.0 0.0 0.0; 0.0 4.0 0.0; 0.0 0.0 4.0]
    @test positions ≈ [0. 0. 0.; 0.5 0.5 0.5]
    @test types == [1, 1]
end

@testset "lattice reduce(export)" begin
    latt = [4.0 20.0 0.0; 0.0 2.0 0.0; 0.0 0.0 12.0]
    niggli_reduce!(latt, 1e-3)
    @test latt ≈ [0.0 -2.0 0.0; 4.0 0.0 0.0; 0.0 0.0 12.0]

    latt = [4.0 20.0 0.0; 0.0 2.0 0.0; 0.0 0.0 12.0]
    delaunay_reduce!(latt, 1e-3)
    @test latt ≈ [0.0 2.0 0.0; -4.0 -0.0 0.0; -0.0 -0.0 12.0]
end

# @testset "reciprocal mesh(export)" begin
#     latt = [5.0 0.0 5.0; 0.0 5.0 5.0; 5.0 5.0 0.0]
#     positions = Array{Float64, 2}([0.0 0.25; 0.0 0.25; 0.0 0.25])
#     types = [1, 1]
#     na = 2
#     mesh = [3,3,3]
#     is_shift = [true,true,true]
#     nir, grid_address, ir_mapping_table =
#         ir_reciprocal_mesh(mesh, is_shift, true, latt, positions, types, na, 1e-5)
#     @test nir == 4
# end

@testset "symmetry(export)" begin
    latt = [4.0 0.0 0.0; 2.0 3.4641 0.0; 0.0 0.0 12.0]
    positions = [0.0 0.0 0.0; 1/3 1/3 1/3]
    types = [1, 1]
    num_atom = 2
    @test get_spacegroup(latt, positions, types, num_atom) == ("P-3m1", 164)

    rots, trans = get_symmetry(latt, positions, types, 1e-3)
    @test length(rots) == length(trans) == 12
    @test size(rots[1]) == (3, 3)
    @test size(trans[1]) == (3,)
end


text:
using Plots

function save_binned_surface_plots(
    x,
    y,
    z,
    title,
    filename,
    n_plots = (3, 3),
    z_label_prefix = "z",
    n_digits = 5,
)
    n_z_partitions = prod(n_plots)
    z_min_global = min(z...)
    z_max_global = max(z...)
    Δz = (z_max_global - z_min_global) / n_z_partitions
    z_min = ntuple(i -> z_min_global + (i - 1) * Δz, n_z_partitions)
    z_max = ntuple(i -> z_min_global + (i) * Δz, n_z_partitions)
    p = []
    for i in 1:n_z_partitions
        mask = z_min[i] .<= z .<= z_max[i]
        x_i = x[mask]
        y_i = y[mask]
        sz_min = string(z_min[i])[1:min(n_digits, length(string(z_min[i])))]
        sz_max = string(z_max[i])[1:min(n_digits, length(string(z_max[i])))]
        p_i = plot(
            x_i,
            y_i,
            title = "$(title), in ($sz_min, $sz_max)",
            seriestype = :scatter,
            markersize = 5,
        )
        push!(p, p_i)
    end
    plot(p..., layout = n_plots, legend = false)
    savefig(filename)
end;


text:
# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# Licensed under ASL - see ASL.md for terms and conditions.
# --------------------------------------------------------------------------

export get_orders, get_nl

import ACE: standardevaluator, graphevaluator
using SparseArrays: SparseMatrixCSC, sparse
using LinearAlgebra: mul!

"""
`struct RPIBasis`
"""
struct RPIBasis{T, BOP, NZ, TIN} <: IPBasis
   pibasis::PIBasis{BOP, NZ, TIN}
   A2Bmaps::NTuple{NZ, SparseMatrixCSC{T, Int}}
   Bz0inds::NTuple{NZ, UnitRange{Int}}
end

Base.length(basis::RPIBasis, iz0::Integer) = size(basis.A2Bmaps[iz0], 1)

Base.length(basis::RPIBasis) = sum(length(basis, iz0)
                                    for iz0 = 1:numz(basis.pibasis))

fltype(::RPIBasis{T}) where {T}  = T

zlist(basis::RPIBasis) = zlist(basis.pibasis)

cutoff(basis::RPIBasis) = cutoff(basis.pibasis)

standardevaluator(basis::RPIBasis) =
      RPIBasis( standardevaluator(basis.pibasis),
                basis.A2Bmaps, basis.Bz0inds )

graphevaluator(basis::RPIBasis) =
      RPIBasis( graphevaluator(basis.pibasis),
                basis.A2Bmaps, basis.Bz0inds )

_basisfcnidx(basis::RPIBasis, iz0::Integer, ib::Integer) = basis.Bz0inds[iz0][ib]

# ------------------------------------------------------------------------
#    Utilities / Basis analysis code
# ------------------------------------------------------------------------

function get_orders(basis::RPIBasis)
   # get the correlation orders of the PIbasis ...
   Ns_pi = zeros(Int, length(basis.pibasis))
   for iz0 = 1:numz(basis)
      inner = basis.pibasis.inner[iz0]
      Ns_pi[inner.AAindices] .= inner.orders
   end
   # ... and use them to construct the correlation orders of the RPI basis
   Ns = zeros(Int, length(basis))
   for iz0 = 1:numz(basis)
      # loop over basis functions belonging to centre-species iz0
      for ib = 1:length(basis, iz0)
         # find one of the indices of the PI basis that the current
         # RPI basis function belongs to; they are all equivalent in
         # terms of correlation-order, so only one of them matters.
         iPI = findfirst(basis.A2Bmaps[iz0][ib,:] .!= 0)
         Ns[_basisfcnidx(basis, iz0, ib)] = Ns_pi[iPI]
      end
   end
   return Ns
end

function get_nl(basis::RPIBasis)
   # get the correlation orders of the PIbasis ...
   nnll_pi = Vector{Vector{Any}}(undef, length(basis.pibasis))
   for iz0 = 1:numz(basis)
      inner = basis.pibasis.inner[iz0]
      for (key, val) in inner.b2iAA
         b = [ (l = b.l, n = b.n) for b in key.oneps ]
         nnll_pi[inner.AAindices[val]] = b
      end
   end
   # ... and use them to construct the correlation orders of the RPI basis
   nnll = Vector{Vector{Any}}(undef, length(basis))
   for iz0 = 1:numz(basis)
      # loop over basis functions belonging to centre-species iz0
      for ib = 1:length(basis, iz0)
         # find one of the indices of the PI basis that the current
         # RPI basis function belongs to; they are all equivalent in
         # terms of correlation-order, so only one of them matters.
         iPI = findfirst(basis.A2Bmaps[iz0][ib,:] .!= 0)
         nnll[_basisfcnidx(basis, iz0, ib)] = nnll_pi[iPI]
      end
   end
   return nnll
end

# ------------------------------------------------------------------------
#    FIO code
# ------------------------------------------------------------------------

==(B1::RPIBasis, B2::RPIBasis) = (B1.pibasis == B2.pibasis)

write_dict(basis::RPIBasis) = Dict(
      "__id__" => "ACE_RPIBasis",
      "__v__"  => "v0_8_2",
      "pibasis" => write_dict(basis.pibasis),
      "A2Bmaps" => write_dict.(basis.A2Bmaps),
      "Bz0inds" => [ [ur.start, ur.stop] for ur in basis.Bz0inds ]
   )

# v0.8.2 onwards
function read_dict(::Val{:ACE_RPIBasis}, ::Val{:v0_8_2}, D::Dict)
   pibasis = read_dict(D["pibasis"])
   A2Bmaps = tuple( read_dict.(D["A2Bmaps"])... )
   Bz0inds = tuple( [ ur[1]:ur[2] for ur in D["Bz0inds"] ]... )
   return RPIBasis(pibasis, A2Bmaps, Bz0inds)
end

# old version
read_dict(::Val{:ACE_RPIBasis}, D::Dict) =
   RPIBasis(read_dict(D["pibasis"]))

read_dict(::Val{:SHIPs_RPIBasis}, D::Dict) =
   read_dict(Val{:ACE_RPIBasis}(), D)

# ------------------------------------------------------------------------
#    Basis construction code
# ------------------------------------------------------------------------

RPIBasis(basis1p::OneParticleBasis, N::Integer,
         D::AbstractDegree, maxdeg::Real, constants=false) =
   RPIBasis(PIBasis(basis1p, N, D, maxdeg; filter = RPIFilter(constants)))

function RPIBasis(pibasis::PIBasis)
   basis1p = pibasis.basis1p

   # construct the cg matrices
   rotc = Rot3DCoeffs()
   A2Bmaps = ntuple(iz0 -> _rpi_A2B_matrix(rotc, pibasis, iz0), numz(pibasis))

   # construct the indices within the B vector to which the A2Bmaps map.
   Bz0inds = UnitRange{Int}[]
   idx0 = 0
   for i = 1:length(A2Bmaps)
      len = size(A2Bmaps[i], 1)
      push!(Bz0inds, (idx0+1):(idx0+len))
      idx0 += len
   end

   return RPIBasis(pibasis, A2Bmaps, tuple(Bz0inds...))
end

# TODO NOW: graphevaluator, standardevaluator

struct RPIFilter
   constants::Bool
end

(f::RPIFilter)(pib::PIBasisFcn{0}) = f.constants
(f::RPIFilter)(pib::PIBasisFcn{1}) = (pib.oneps[1].l == 0)
(f::RPIFilter)(pib::PIBasisFcn) = (
      iseven( sum(b.l for b in pib.oneps) ) &&
      (sum(b.m for b in pib.oneps) == 0) )

# _rpi_filter(pib::PIBasisFcn{0}) = false
# _rpi_filter(pib::PIBasisFcn{1}) = (pib.oneps[1].l == 0)
# _rpi_filter(pib::PIBasisFcn) = (
#       iseven( sum(b.l for b in pib.oneps) ) &&
#       (sum(b.m for b in pib.oneps) == 0) )

function _rpi_A2B_matrix(rotc::Rot3DCoeffs,
                         pibasis::PIBasis,
                         iz0)
   # allocate triplet format
   Irow, Jcol, vals = Int[], Int[], real(fltype(pibasis.basis1p))[]
   # count the number of PI basis functions = number of rows
   idxB = 0
   # loop through all (zz, kk, ll) tuples; each specifies 1 to several B
   for i = 1:length(pibasis.inner[iz0])
      # get the specification of the ith basis function
      pib = get_basis_spec(pibasis, iz0, i)
      # skip it unless all m are zero, because we want to consider each
      # (nn, ll) block only once.
      if !all(b.m == 0 for b in pib.oneps)
         continue
      end
      # get the rotation-coefficients for this basis group
      # the bs are the basis functions corresponding to the columns
      U, bcols = _rpi_coupling_coeffs(pibasis, rotc, pib)
      # loop over the rows of Ull -> each specifies a basis function
      for irow = 1:size(U, 1)
         idxB += 1
         # loop over the columns of U / over brows
         for (icol, bcol) in enumerate(bcols)
            # this is a subtle step: bcol and bcol_ordered are equivalent
            # permutation-invariant basis functions. This means we will
            # add the same PI basis function several times, but in the call to
            # `sparse` the values will just be added.
            bcol_ordered = ACE._get_ordered(pibasis, bcol)
            idxAA = pibasis.inner[iz0].b2iAA[bcol_ordered]
            push!(Irow, idxB)
            push!(Jcol, idxAA)
            push!(vals, U[irow, icol])
         end
      end
   end
   # create CSC: [   triplet    ]  nrows   ncols
   return sparse(Irow, Jcol, vals, idxB, length(pibasis.inner[iz0]))
end

# U, bcols = rpi_coupling_coeffs(rotc, pib)

"""
this is essentially a wrapper function around Rotations3D.rpi_basis,
and is just meant to translate between different representations
"""
function _rpi_coupling_coeffs(pibasis, rotc::Rot3DCoeffs, pib::PIBasisFcn{N}
                              ) where {N}
   # convert to zz, ll, nn tuples
   zz, nn, ll, _ = _b2znlms(pib)
   # construct the RPI coupling coefficients
   U, Ms = Rotations3D.rpi_basis(rotc, zz, nn, ll)
   # convert the Ms into basis functions
   rpibs = [ _znlms2b(zz, nn, ll, mm, pib.z0) for mm in Ms ]
   return U, rpibs
end

_rpi_coupling_coeffs(pibasis, rotc::Rot3DCoeffs, pib::PIBasisFcn{0}) =
      [ 1.0 ], []

_b2znlms(pib::PIBasisFcn{N}) where {N} = (
   SVector(ntuple(n -> pib.oneps[n].z, N)...),
   SVector(ntuple(n -> pib.oneps[n].n, N)...),
   SVector(ntuple(n -> pib.oneps[n].l, N)...),
   SVector(ntuple(n -> pib.oneps[n].m, N)...) )

_znlms2b(zz, nn, ll, mm = zero(ll), z0 = AtomicNumber(0)) =
   PIBasisFcn( z0, ntuple(i -> PSH1pBasisFcn(nn[i], ll[i], mm[i], zz[i]),
                          length(zz)) )

function combine(basis::RPIBasis, coeffs)
   picoeffs = ntuple(iz0 -> (coeffs[basis.Bz0inds[iz0]]' * basis.A2Bmaps[iz0])[:],
                     numz(basis.pibasis))
   return PIPotential(basis.pibasis, picoeffs)
end

"""
`scaling(basis::RPIBasis, p[; a2b = abs2, fin = ...])` : same as usual but adds a keyword
argument `use` :
* `a2b = identity` : the old buggy scaling
* `a2b = abs` : the bugfix which keeps the behaviour same/similar
* `a2b = abs2` : a different kind of scaling which I believe is the right one
but behaves a bit differently especially for high body-orders.

The `fin` keyword argument is the "finishing" map. E.g. if we take `a2b = abs2`,
then we should apply a sqrt at the end. This is done automatically here.
"""
function scaling(basis::RPIBasis, p; a2b = abs2,
                 fin = (a2b == abs2 ? sqrt : identity) )
   wwpi = scaling(basis.pibasis, p)
   wwrpi = zeros(Float64, length(basis))
   for iz0 = 1:numz(basis)
      wwpi_iz0 = wwpi[basis.pibasis.inner[iz0].AAindices]
      wwrpi[basis.Bz0inds[iz0]] = a2b.(basis.A2Bmaps[iz0]) * a2b.(wwpi_iz0)
   end
   return fin.(wwrpi)
end

# ------------------------------------------------------------------------
#    Evaluation code
# ------------------------------------------------------------------------

alloc_temp(basis::RPIBasis, args...) =
   ( AA = site_alloc_B(basis.pibasis, args...),
     AAr = real(site_alloc_B(basis.pibasis, args...)),
     tmp_pibasis = alloc_temp(basis.pibasis, args...)
   )

function evaluate!(B, tmp, basis::RPIBasis, Rs, Zs, z0)
   iz0 = z2i(basis, z0)
   AA = site_evaluate!(tmp.AA, tmp.tmp_pibasis, basis.pibasis, Rs, Zs, z0)
   # TODO: this could be done better maybe by adding the real function into
   #       site_evaluate!, or by writing a real version of it...
   AAr = @view tmp.AAr[1:length(AA)]
   AAr[:] .= real.(AA)
   Bview = @view B[basis.Bz0inds[iz0]]
   mul!(Bview, basis.A2Bmaps[iz0], AAr)
   return B
end

# ------- gradient

alloc_temp_d(basis::RPIBasis, Rs::AbstractVector, args...) =
   alloc_temp_d(basis, length(Rs))

alloc_temp_d(basis::RPIBasis, nmax::Integer) =
    (
    AA = site_alloc_B(basis.pibasis),
    AAr = real(site_alloc_B(basis.pibasis)),
    tmp_pibasis = alloc_temp(basis.pibasis, nmax),
    dAA = site_alloc_dB(basis.pibasis, nmax),
    tmpd_pibasis = alloc_temp_d(basis.pibasis, nmax),
    )

# TODO: evaluate also B??? the interface seems to command it.
function evaluate_d!(B, dB, tmpd, basis::RPIBasis, Rs, Zs, z0)
   iz0 = z2i(basis, z0)
   AA, dAA = tmpd.AA, tmpd.dAA
   site_evaluate_d!(AA, dAA, tmpd.tmpd_pibasis, basis.pibasis, Rs, Zs, z0)
   len = length(basis.pibasis.inner[iz0])
   for i = 1:len
      AA[i] = real(AA[i])
   end
   mul!((@view B[basis.Bz0inds[iz0]]), basis.A2Bmaps[iz0], (@view AA[1:len]))
   for j = 1:length(Rs)
      # ∂∏A / ∂𝐫ⱼ
      dAAj = @view dAA[1:len, j]
      for i = 1:length(dAAj)
         # TODO: we should incorporate the `real` operation into the PIBasis?
         @inbounds dAAj[i] = real.(dAAj[i])
      end
      # copy into B
      dBview = @view dB[basis.Bz0inds[iz0], j]
      mul!(dBview, basis.A2Bmaps[iz0], dAAj)
   end
   return dB
end


text:
using LinearMaps
using IterativeSolvers
using Statistics
import Base: @kwdef

# Mixing rules: (ρin, ρout) => ρnext, where ρout is produced by diagonalizing the
# Hamiltonian at ρin These define the basic fix-point iteration, that are then combined with
# acceleration methods (eg anderson). For the mixing interface we use `δF = ρout - ρin` and
# `δρ = ρnext - ρin`, such that the mixing interface is
# `mix_density(mixing, basis, δF; kwargs...) -> δρ` with the user being assumed to add this
# to ρin to get ρnext. All these methods attempt to approximate the inverse Jacobian of the
# SCF step, ``J^-1 = (1 - χ0 (vc + K_{xc}))^-1``, where vc is the Coulomb and ``K_{xc}`` the
# exchange-correlation kernel. Note that "mixing" is sometimes used to refer to the combined
# process of formulating the fixed-point and solving it; we call "mixing" only the first part
# The notation in this file follows Herbst, Levitt arXiv:2009.01665

# Mixing can be done in the potential or the density. By default we assume
# the dielectric model is so simple that both types of mixing are identical.
# If mixing is done in the potential, the interface is
# `mix_potential(mixing, basis, δF; kwargs...) -> δV`
abstract type Mixing end
function mix_potential(args...; kwargs...)
    mix_density(args...; kwargs...)
end

@doc raw"""
Simple mixing: ``J^{-1} ≈ 1``
"""
struct SimpleMixing <: Mixing; end
mix_density(::SimpleMixing, ::PlaneWaveBasis, δF; kwargs...) = δF

@doc raw"""
Kerker mixing: ``J^{-1} ≈ \frac{|G|^2}{k_{TF}^2 + |G|^2}``
where ``k_{TF}`` is the Thomas-Fermi wave vector. For spin-polarized calculations
by default the spin density is not preconditioned. Unless a non-default value
for ``ΔDOS_Ω`` is specified. This value should roughly be the expected difference in density
of states (per unit volume) between spin-up and spin-down.

Notes:
- Abinit calls ``1/k_{TF}`` the dielectric screening length (parameter *dielng*)
"""
@kwdef struct KerkerMixing <: Mixing
    # Default kTF parameter suggested by Kresse, Furthmüller 1996 (kTF=1.5Å⁻¹)
    # DOI 10.1103/PhysRevB.54.11169
    kTF::Real    = 0.8  # == sqrt(4π (DOS_α + DOS_β)) / Ω
    ΔDOS_Ω::Real = 0.0  # == (DOS_α - DOS_β) / Ω
end

@timing "KerkerMixing" function mix_density(mixing::KerkerMixing, basis::PlaneWaveBasis,
                                            δF; kwargs...)
    T      = eltype(δF)
    G²     = [sum(abs2, G) for G in G_vectors_cart(basis)]
    kTF    = T.(mixing.kTF)
    ΔDOS_Ω = T.(mixing.ΔDOS_Ω)

    # TODO This can be improved to use less copies for the new (α, β) interface

    # For Kerker the model dielectric written as a 2×2 matrix in spin components is
    #     1 - [-DOSα      0] * [1 1]
    #         [    0  -DOSβ]   [1 1] * (4π/G²)
    # which maps (δρα, δρβ)ᵀ to (δFα, δFβ)ᵀ and where DOSα and DOSβ is the density
    # of states per unit volume in the spin-up and spin-down channels. After basis
    # transformation to a mapping (δρtot, δρspin)ᵀ to (δFtot, δFspin)ᵀ this becomes
    #     [(G² + kTF²)    0]
    #     [ 4π * ΔDOS    G²] / G²
    # where we defined kTF² = 4π * (DOSα + DOSβ) and ΔDOS = DOSα - DOSβ.
    # Gaussian elimination on this matrix yields for the linear system ε δρ = δF
    #     δρtot  = G² δFtot / (G² + kTF²)
    #     δρspin = δFspin - 4π * ΔDOS / (G² + kTF²) δFtot

    δF_fourier     = r_to_G(basis, δF)
    δFtot_fourier  = total_density(δF_fourier)
    δFspin_fourier = spin_density(δF_fourier)

    δρtot_fourier = δFtot_fourier .* G² ./ (kTF.^2 .+ G²)
    δρtot = G_to_r(basis, δρtot_fourier)

    # Copy DC component, otherwise it never gets updated
    δρtot .+= mean(total_density(δF)) .- mean(δρtot)

    if basis.model.n_spin_components == 1
        ρ_from_total_and_spin(δρtot, nothing)
    else
        δρspin_fourier = @. δFspin_fourier - δFtot_fourier * (4π * ΔDOS_Ω) / (kTF^2 + G²)
        δρspin = G_to_r(basis, δρspin_fourier)
        ρ_from_total_and_spin(δρtot, δρspin)
    end
end

@doc raw"""
The same as [`KerkerMixing`](@ref), but the Thomas-Fermi wavevector is computed
from the current density of states at the Fermi level.
"""
@kwdef struct KerkerDosMixing <: Mixing
    adjust_temperature = IncreaseMixingTemperature()
end
@timing "KerkerDosMixing" function mix_density(mixing::KerkerDosMixing, basis::PlaneWaveBasis,
                                               δF; εF, eigenvalues, kwargs...)
    if iszero(basis.model.temperature)
        return mix_density(SimpleMixing(), basis, δF)
    else
        n_spin = basis.model.n_spin_components
        Ω = basis.model.unit_cell_volume
        temperature = mixing.adjust_temperature(basis.model.temperature; kwargs...)
        dos_per_vol  = compute_dos(εF, basis, eigenvalues; temperature) ./ Ω
        kTF  = sqrt(4π * sum(dos_per_vol))
        ΔDOS_Ω = n_spin == 2 ? dos_per_vol[1] - dos_per_vol[2] : 0.0
        mix_density(KerkerMixing(kTF=kTF, ΔDOS_Ω=ΔDOS_Ω), basis, δF)
    end
end

@doc raw"""
We use a simplification of the Resta model DOI 10.1103/physrevb.16.2717 and set
``χ_0(q) = \frac{C_0 G^2}{4π (1 - C_0 G^2 / k_{TF}^2)}``
where ``C_0 = 1 - ε_r`` with ``ε_r`` being the macroscopic relative permittivity.
We neglect ``K_\text{xc}``, such that
``J^{-1} ≈ \frac{k_{TF}^2 - C_0 G^2}{ε_r k_{TF}^2 - C_0 G^2}``

By default it assumes a relative permittivity of 10 (similar to Silicon).
`εr == 1` is equal to `SimpleMixing` and `εr == Inf` to `KerkerMixing`.
The mixing is applied to ``ρ`` and ``ρ_\text{spin}`` in the same way.
"""
@kwdef struct DielectricMixing <: Mixing
    kTF::Real = 0.8
    εr::Real  = 10
end
@timing "DielectricMixing" function mix_density(mixing::DielectricMixing, basis::PlaneWaveBasis,
                                                δF; kwargs...)
    T = eltype(δF)
    εr = T(mixing.εr)
    kTF = T(mixing.kTF)
    εr == 1               && return mix_density(SimpleMixing(), basis, δF)
    εr > 1 / sqrt(eps(T)) && return mix_density(KerkerMixing(kTF=kTF), basis, δF)

    C0 = 1 - εr
    Gsq = [sum(abs2, G) for G in G_vectors_cart(basis)]
    δF_fourier = r_to_G(basis, δF)
    δρ = @. δF_fourier * (kTF^2 - C0 * Gsq) / (εr * kTF^2 - C0 * Gsq)
    δρ = G_to_r(basis, δρ)
    δρ .+= mean(δF) .- mean(δρ)
end

@doc raw"""
The model for the susceptibility is
```math
\begin{aligned}
    χ_0(r, r') &= (-D_\text{loc}(r) δ(r, r') + D_\text{loc}(r) D_\text{loc}(r') / D) \\
    &+ \sqrt{L(x)} \text{IFFT} \frac{C_0 G^2}{4π (1 - C_0 G^2 / k_{TF}^2)} \text{FFT} \sqrt{L(x)}
\end{aligned}
```
where ``C_0 = 1 - ε_r``, ``D_\text{loc}`` is the local density of states,
``D`` is the density of states and
the same convention for parameters are used as in [`DielectricMixing`](@ref).
Additionally there is the real-space localization function `L(r)`.
For details see Herbst, Levitt 2020 arXiv:2009.01665

Important `kwargs` passed on to [`χ0Mixing`](@ref)
- `RPA`: Is the random-phase approximation used for the kernel (i.e. only Hartree kernel is
  used and not XC kernel)
- `verbose`: Run the GMRES in verbose mode.
- `reltol`: Relative tolerance for GMRES
"""
function HybridMixing(;εr=1.0, kTF=0.8, localization=identity,
                      adjust_temperature=IncreaseMixingTemperature(), kwargs...)
    χ0terms = [DielectricModel(εr=εr, kTF=kTF, localization=localization),
               LdosModel(;adjust_temperature)]
    χ0Mixing(; χ0terms=χ0terms, kwargs...)
end

@doc raw"""
The model for the susceptibility is
```math
\begin{aligned}
    χ_0(r, r') &= (-D_\text{loc}(r) δ(r, r') + D_\text{loc}(r) D_\text{loc}(r') / D)
\end{aligned}
```
where ``D_\text{loc}`` is the local density of states,
``D`` is the density of states. For details see Herbst, Levitt 2020 arXiv:2009.01665.

Important `kwargs` passed on to [`χ0Mixing`](@ref)
- `RPA`: Is the random-phase approximation used for the kernel (i.e. only Hartree kernel is
  used and not XC kernel)
- `verbose`: Run the GMRES in verbose mode.
- `reltol`: Relative tolerance for GMRES
"""
function LdosMixing(; adjust_temperature=IncreaseMixingTemperature(), kwargs...)
    χ0Mixing(; χ0terms=[LdosModel(;adjust_temperature)], kwargs...)
end

@doc raw"""
Generic mixing function using a model for the susceptibility composed of the sum of the `χ0terms`.
For valid `χ0terms` See the subtypes of `χ0Model`. The dielectric model is solved in
real space using a GMRES. Either the full kernel (`RPA=false`) or only the Hartree kernel
(`RPA=true`) are employed. `verbose=true` lets the GMRES run in verbose mode
(useful for debugging).
"""
@kwdef struct χ0Mixing <: Mixing
    RPA::Bool = true       # Use RPA, i.e. only apply the Hartree and not the XC Kernel
    χ0terms   = χ0Model[Applyχ0Model()]  # The terms to use as the model for χ0
    verbose::Bool = false   # Run the GMRES verbosely
    reltol::Float64 = 0.01  # Relative tolerance for GMRES
end

@views @timing "χ0Mixing" function mix_density(mixing::χ0Mixing, basis, δF::AbstractArray{T};
                                               ρin, kwargs...) where {T}
    # Initialise χ0terms and remove nothings (terms that don't yield a contribution)
    χ0applies = filter(!isnothing, [χ₀(basis; ρin=ρin, kwargs...) for χ₀ in mixing.χ0terms])

    # If no applies left, do not bother running GMRES and directly do simple mixing
    isempty(χ0applies) && return mix_density(SimpleMixing(), basis, δF)

    # Solve (ε^†) δρ = δF with ε^† = (1 - χ₀ vc) and χ₀ given as the sum of the χ0terms
    devec(x) = reshape(x, size(δF))
    function dielectric_adjoint(δF)
        δF = devec(δF)
        # Apply Kernel (just vc for RPA and (vc + K_{xc}) if not RPA)
        δV = apply_kernel(basis, δF; ρ=ρin, RPA=mixing.RPA)
        δV .-= mean(δV)
        εδF = copy(δF)
        for apply_term! in χ0applies
            apply_term!(εδF, δV, -1)  # εδF .-= χ₀ * δV
        end
        εδF .-= mean(εδF)
        vec(εδF)
    end

    DC_δF = mean(δF)
    δF .-= DC_δF
    ε  = LinearMap{T}(dielectric_adjoint, length(δF))
    δρ = devec(gmres(ε, vec(δF), verbose=mixing.verbose, reltol=T(mixing.reltol)))
    δρ .+= DC_δF  # Set DC from δF
    δρ
end

@timing "χ0Mixing" function mix_potential(mixing::Mixing, basis::χ0Mixing, δF::AbstractArray; kwargs...)
    error("Not yet implemented.")
end

"""
Increase the temperature used for computing the SCF preconditioners. Initially the temperature
is increased by a `factor`, which is then smoothly lowered towards the temperature used
within the model as the SCF converges. Once the density change is below `above_ρdiff` the
mixing temperature is equal to the model temperature.
"""
function IncreaseMixingTemperature(;factor=25, above_ρdiff=1e-2, temperature_max=0.5)
    function callback(temperature; n_iter, ρin=nothing, ρout=nothing, info...)
        if iszero(temperature) || temperature > temperature_max
            return temperature
        elseif isnothing(ρin) || isnothing(ρout)
            return temperature
        elseif n_iter ≤ 1
            return factor * temperature
        end

        # Continuous piecewise linear function on a logarithmic scale
        # In [log(above_ρdiff), log(above_ρdiff) + 1] it switches from 1 to factor
        ρdiff = norm(ρout .- ρin)
        enhancement = clamp(1 + (factor - 1) * log10(ρdiff / above_ρdiff), 1, factor)

        # Between SCF iterations temperature may never grow
        temperature = clamp(enhancement * temperature, temperature, temperature_max)
        temperature_max = temperature
        return temperature
    end
end


text:
using FEASTSolver
using LinearAlgebra
using MatrixMarket

A0 = mmread("data/system5A0.mtx")
A1 = mmread("data/system5A1.mtx")
A2 = mmread("data/system5A2.mtx")

function T(z::ComplexF64)
	return (A2 .* (z*z)) .+ (A1 .* z) .+ A0
end

R = 0.05
# C = complex(-0.8, 0.8)
C = complex(-1.55, 0.0)

# e, v, res = nlfeast!(T, rand(ComplexF64,1000,3), 2^2, 0, c=complex(-1.55,0.0), r=0.05)
# e, v, res = @timev nlfeast!(T, rand(ComplexF64,1000,30), 2^3, 10, c=complex(-1.55,0.0), r=0.05)
e, v, res = nlfeast_moments!(T, rand(ComplexF64,1000,10), 2^2, 0, c=C, r=R, debug=false, moments=3)
e, v, res = @timev nlfeast_moments!(T, rand(ComplexF64,1000,80), 2^5, 10, c=C, r=R, debug=true, moments=2, ϵ=10e-16, store=true)
# e, v, res = beyn(T, A0, rand(ComplexF64,1000,120), 2^9; c=complex(-1.55,0.0), r=0.05)
# display(e)
# display(res)
inside(x) = in_contour(x, C, R)
# print("\nmax res inside: ")
# println(maximum(res[inside.(e)]))
# print("number inside : ")
# println(size(e[inside.(e)])[1])
# # display(e)
# display(res[inside.(e)])

in_eig = e[inside.(e)]
in_res = res[inside.(e)]
in_res_conv = in_res[in_res .<= 1e-3]
in_eig_conv = in_eig[in_res .<= 1e-3]
print("\nmax res inside: ")
println(maximum(res[inside.(e)]))
print("max res inside non spurious: ")
println(maximum(in_res_conv))
print("number inside : ")
println(size(e[inside.(e)])[1])
print("number inside converged : ")
println(size(in_eig_conv)[1])
# display(in_eig_conv)


text:
# Tests for the AdaBoost ensemble
function t_adaboost()

	# M1
 	Test.@test try
		L = 10 # ensemble size

		# Prepare data
		D=j4pr.DataGenerator.fish(20);
		(tr,ts)=j4pr.splitobs(j4pr.shuffleobs(D),0.3);
		yorig=-D;
		yu=sort(unique(yorig));
		ytr=Int.(j4pr.ohenc_integer(-tr, yu)); # knn works only with integer labels

		# Construct execution function (has to return a vector of labels)
		f_exec=(m,x)->j4pr.targets(indmax, j4pr.DecisionStump.stump_exec(m,j4pr.getobs(x)))

		# Train
		ae=j4pr.AdaBoost.adaboost_train(+tr, ytr, L, x->j4pr.DecisionStump.stump_train(x[1],x[2]), f_exec, j4pr.AdaBoost.AdaBoostM1())

		# Execute
		j4pr.AdaBoost.adaboost_exec(ae,+D)
		true
	catch
		false
	end

	# M2
 	Test.@test try
		L = 10 # ensemble size

		# Prepare data
		D=j4pr.DataGenerator.fish(20);
		(tr,ts)=j4pr.splitobs(j4pr.shuffleobs(D),0.3);
		yorig=-D;
		yu=sort(unique(yorig));
		ytr=Int.(j4pr.ohenc_integer(-tr, yu));

		# Construct execution function (has to return a matrix of observation probabilities)
		f_exec=(m,x)->j4pr.DecisionStump.stump_exec(m,j4pr.getobs(x))

		# Train
		ae=j4pr.AdaBoost.adaboost_train(+tr, ytr, L, x->j4pr.DecisionStump.stump_train(x[1],x[2]), f_exec, j4pr.AdaBoost.AdaBoostM2())

		# Execute
		j4pr.AdaBoost.adaboost_exec(ae,+D)
		true
	catch
		false
	end

	# AdaBoost FunctionCell interface
	A = j4pr.DataGenerator.fish(20) 			# load data
	tr,ts = j4pr.splitobs(j4pr.shuffleobs(A), 0.5)		# shuffle and split data
	w=j4pr.parzen(0.1) 					# define classifier
	L=10

	Test.@test try
		m1 = j4pr.adaboost(w, L; boost_type=j4pr.AdaBoost.AdaBoostM1())
		m2 = j4pr.adaboost(w, L; boost_type=j4pr.AdaBoost.AdaBoostM2())

		# Train
		m1t = m1(tr)
		m2t = m2(tr)

		# Execute
		m1t(ts);
		m2t(ts);

		# If we're here, everything worked
		true

	catch
		false
	end

end


text:
################################################################################
#
#  NfRel/NfRel_ns.jl : non-simple relative fields
#
# This file is part of Hecke.
#
# Copyright (c) 2015, 2016, 2017: Claus Fieker, Tommy Hofmann
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#
#  Copyright (C) 2017 Tommy Hofmann, Claus Fieker
#
################################################################################

export NfRel_ns, simple_extension

#= trivial example
Qx, x = PolynomialRing(FlintQQ)
QQ = number_field(x-1)[1]
QQt, t = QQ["t"]
K, gK = number_field([t^2-2, t^2-3, t^2-5, t^2-7])
[ minpoly(x) for x = gK]
S, mS = Hecke.simple_extension(K)
mS(gen(S))
mS\gens(K)[2]
=#

#to make the MPoly module happy, divrem needs it...
function Nemo.div(a::nf_elem, b::nf_elem)
  return a//b
end

function Nemo.rem(a::nf_elem, b::nf_elem)
  return parent(a)(0)
end

#non-simple fields are quotients by multivariate polynomials
#this could be extended to arbitrary zero-dimensional quotients, but
#I don't need this here.

#mostly copied from NfRel I am afraid..

################################################################################
#
#  Copy
#
################################################################################

function Base.deepcopy_internal(a::NfRel_nsElem{T}, dict::IdDict) where T
  z = NfRel_nsElem{T}(Base.deepcopy_internal(data(a), dict))
  z.parent = parent(a)
  return z
end

#julia's a^i needs copy
function Base.copy(a::NfRel_nsElem)
  return parent(a)(a.data)
end

################################################################################
#
#  Comply with Nemo ring interface
#
################################################################################

Nemo.elem_type(::Type{NfRel_ns{T}}) where {T} = NfRel_nsElem{T}

Nemo.elem_type(::NfRel_ns{T}) where {T} = NfRel_nsElem{T}

Nemo.parent_type(::Type{NfRel_nsElem{T}}) where {T} = NfRel_ns{T}

order_type(K::NfRel_ns{T}) where {T} = NfRelOrd{T, frac_ideal_type(order_type(base_field(K)))}

order_type(::Type{NfRel_ns{T}}) where {T} = NfRelOrd{T, frac_ideal_type(order_type(parent_type(T)))}

Nemo.needs_parentheses(::NfRel_nsElem) = true

Nemo.isnegative(x::NfRel_nsElem) = Nemo.isnegative(data(x))

Nemo.show_minus_one(::Type{NfRel_nsElem{T}}) where {T} = true

function Nemo.iszero(a::NfRel_nsElem)
  reduce!(a)
  return iszero(data(a))
end

function Nemo.isone(a::NfRel_nsElem)
  reduce!(a)
  return isone(data(a))
end

Nemo.zero(K::NfRel_ns) = K(Nemo.zero(parent(K.pol[1])))

Nemo.one(K::NfRel_ns) = K(Nemo.one(parent(K.pol[1])))

Nemo.one(a::NfRel_nsElem) = one(a.parent)

(K::NfRel_ns{T})(a::NfRel_nsElem) where {T} = a

################################################################################
#
#  Promotion
#
################################################################################

Nemo.promote_rule(::Type{NfRel_nsElem{S}}, ::Type{T}) where {T <: Integer, S} = NfRel_nsElem{S}

Nemo.promote_rule(::Type{NfRel_nsElem{T}}, ::Type{fmpz}) where {T <: Nemo.RingElement} = NfRel_nsElem{T}

Nemo.promote_rule(::Type{NfRel_nsElem{T}}, ::Type{fmpq}) where {T <: Nemo.RingElement} = NfRel_nsElem{T}

Nemo.promote_rule(::Type{NfRel_nsElem{T}}, ::Type{T}) where {T} = NfRel_nsElem{T}

Nemo.promote_rule(::Type{NfRel_nsElem{T}}, ::Type{NfRel_nsElem{T}}) where T <: Nemo.RingElement = NfRel_nsElem{T}

function Nemo.promote_rule(::Type{NfRel_nsElem{T}}, ::Type{U}) where {T <: Nemo.RingElement, U <: Nemo.RingElement}
  Nemo.promote_rule(T, U) == T ? NfRel_nsElem{T} : Union{}
end

################################################################################
#
#  Field access
#
################################################################################

@inline base_field(a::NfRel_ns{T}) where {T} = a.base_ring::parent_type(T)

@inline Nemo.data(a::NfRel_nsElem) = a.data

@inline Nemo.parent(a::NfRel_nsElem{T}) where {T} = a.parent::NfRel_ns{T}

issimple(a::NfRel_ns) = false

################################################################################
#
#  Reduction
#
################################################################################

function reduce!(a::NfRel_nsElem)
  q, a.data = divrem(a.data, parent(a).pol)
  return a
end

################################################################################
#
#  String I/O
#
################################################################################

function Base.show(io::IO, a::NfRel_ns)
  print(io, "non-simple Relative number field over\n")
  print(io, base_field(a), "\n")
  print(io, " with defining polynomials ", a.pol)
end

#TODO: this is a terrible show func.
function Base.show(io::IO, a::NfRel_nsElem)
  f = data(a)
  AbstractAlgebra.Generic._show(io, f, [string(s) for s = a.parent.S])
  return nothing
end

################################################################################
#
#  Constructors and parent object overloading
#
################################################################################

function Nemo.NumberField(f::Array{Generic.Poly{T}, 1}, s::String="_\$"; cached::Bool = false, check::Bool = true) where T
  S = Symbol(s)
  R = base_ring(f[1])
  Rx, x = PolynomialRing(R, length(f), s)
  K = NfRel_ns(f, [f[i](x[i]) for i=1:length(f)], [Symbol("$s$i") for i=1:length(f)])
  if check
    if !check_consistency(K)
      error("The fields are not linearly disjoint!")
    end
  end
  return K, gens(K)
end

Nemo.gens(K::NfRel_ns) = [K(x) for x = gens(parent(K.pol[1]))]

function (K::NfRel_ns{T})(a::Generic.MPoly{T}) where T
  q, w = divrem(a, K.pol)
  z = NfRel_nsElem{T}(w)
  z.parent = K
  return z
end

function (K::NfRel_ns{T})(a::T) where T
  parent(a) != base_ring(parent(K.pol[1])) == error("Cannot coerce")
  z = NfRel_nsElem{T}(parent(K.pol[1])(a))
  z.parent = K
  return z
end

function (K::NfRel_ns{T})(a::NfRel_nsElem{T}) where T
  parent(a) != K == error("Cannot coerce")
  return a
end

(K::NfRel_ns)(a::Integer) = K(parent(K.pol[1])(a))

(K::NfRel_ns)(a::Rational{T}) where {T <: Integer} = K(parent(K.pol[1])(a))

(K::NfRel_ns)(a::fmpz) = K(parent(K.pol[1])(a))

(K::NfRel_ns)(a::fmpq) = K(parent(K.pol[1])(a))

(K::NfRel_ns)() = zero(K)

Nemo.gen(K::NfRel_ns) = K(Nemo.gen(parent(K.pol[1])))

################################################################################
#
#  Unary operators
#
################################################################################

function Base.:(-)(a::NfRel_nsElem)
  return parent(a)(-data(a))
end

################################################################################
#
#  Binary operators
#
################################################################################

function Base.:(+)(a::NfRel_nsElem{T}, b::NfRel_nsElem{T}) where {T}
  return parent(a)(data(a) + data(b))
end

function Base.:(-)(a::NfRel_nsElem{T}, b::NfRel_nsElem{T}) where {T}
  return parent(a)(data(a) - data(b))
end

function Base.:(*)(a::NfRel_nsElem{T}, b::NfRel_nsElem{T}) where {T}
  return parent(a)(data(a) * data(b))
end

function Base.:(//)(a::NfRel_nsElem{T}, b::NfRel_nsElem{T}) where {T}
  return div(a, b)
end

function Nemo.div(a::NfRel_nsElem{T}, b::NfRel_nsElem{T}) where {T}
  return a*inv(b)
end

Nemo.divexact(a::NfRel_nsElem, b::NfRel_nsElem) = div(a, b)
################################################################################
#
#  Powering
#
################################################################################
#via julia

function Base.:(^)(a::NfRel_nsElem{T}, b::Integer) where T
  if b < 0
    return inv(a)^(-b)
  elseif b == 0
    return one(parent(a))
  elseif b == 1
    return deepcopy(a)
  elseif mod(b, 2) == 0
    c = a^(div(b, 2))
    return c*c
  else mod(b, 2) == 1
    return a^(b - 1)*a
  end
end

function Base.:(^)(a::NfRel_nsElem{T}, b::fmpz) where T
  if b < 0
    return inv(a)^(-b)
  elseif b == 0
    return parent(a)(1)
  elseif b == 1
    return deepcopy(a)
  elseif mod(b, 2) == 0
    c = a^(div(b, 2))
    return c*c
  else mod(b, 2) == 1
    return a^(b - 1)*a
  end
end

################################################################################
#
#  Comparison
#
################################################################################

function Base.:(==)(a::NfRel_nsElem{T}, b::NfRel_nsElem{T}) where T
  reduce!(a)
  reduce!(b)
  return data(a) == data(b)
end

################################################################################
#
#  Unsafe operations
#
################################################################################

function Nemo.mul!(c::NfRel_nsElem{T}, a::NfRel_nsElem{T}, b::NfRel_nsElem{T}) where {T}
  mul!(c.data, a.data, b.data)
  c = reduce!(c)
  return c
end

function Nemo.addeq!(b::NfRel_nsElem{T}, a::NfRel_nsElem{T}) where {T}
  addeq!(b.data, a.data)
  b = reduce!(b)
  return b
end

function Base.hash(a::NfRel_nsElem{nf_elem}, b::UInt)
  reduce!(a)
  return hash(a.data, b)
end

###############################################################################
# other stuff, trivia and non-trivia
###############################################################################

function Nemo.degree(K::NfRel_ns)
  return prod([total_degree(x) for x=K.pol])
end

function (R::Generic.PolyRing{nf_elem})(f::Generic.MPoly)
  if length(f)==0
    return R()
  end
  j = 1
  c = 0
  while j<= ngens(parent(f))
    if f.exps[j, 1] != 0
      if c==0
        c = j
      else
        error("poly is not univariate")
      end
    end
    j += 1
  end
  g = R()
  for i=1:length(f)
    setcoeff!(g, Int(f.exps[c, i]), f.coeffs[i])
  end
  return g
end

#non-optimal...
function basis(K::NfRel_ns)
  b = NfRel_nsElem[]
  g = gens(K)
  for i=CartesianIndices(Tuple(1:total_degree(f) for f = K.pol))
    push!(b, prod(g[j]^(i[j]-1) for j=1:length(i)))
  end
  return b
end

function basis_matrix(a::Vector{NfRel_nsElem{T}}) where {T <: NumFieldElem}
  @assert length(a) > 0
  K = parent(a[1])
  M = zero_matrix(base_field(K), length(a), degree(K))
  for i in 1:length(a)
    elem_to_mat_row!(M, i, a[i])
  end
  return M
end

function elem_to_mat_row!(M::Generic.Mat{T}, i::Int, a::NfRel_nsElem{T}) where T
  a.parent
  K = parent(a)
  #C = CartesianIndices(Tuple(0:total_degree(f)-1 for f = K.pol))
  #L = LinearIndices(C)
  #CC = [UInt[c[i] for i=length(K.pol):-1:1] for c = C]
  for j=1:ncols(M)
    M[i, j] = zero(base_field(K))
  end
  for j=1:length(a.data)
    #p = findfirst(isequal(a.data.exps[:, j]), CC)
    #@assert p !== nothing
    M[i, monomial_to_index(j, a)] = a.data.coeffs[j]
  end
end

function elem_from_mat_row(K::NfRel_ns{T}, M::Generic.Mat{T}, i::Int) where T
  a = K()
  t = K()
  b = basis(K)
  for c = 1:ncols(M)
    a.parent = K
    b[c].parent = K
    a = a + M[i, c]*b[c]
    b[c].parent
  end
  a.parent
  return a
end

function monomial_to_index(i::Int, a::NfRel_nsElem)
  K = parent(a)
  n = ngens(K)
  idx = a.data.exps[1, i]
  for j=2:n
    idx *= total_degree(K.pol[n - j + 1])
    idx += a.data.exps[j, i]
  end
  return Int(idx+1)
end

function SRow(a::NfRel_nsElem)
  sr = SRow(base_field(parent(a)))
  for i=1:length(a.data)
    push!(sr.pos, monomial_to_index(i, a))
    push!(sr.values, a.data.coeffs[i])
  end
  p = sortperm(sr.pos)
  sr.pos = sr.pos[p]
  sr.values = sr.values[p]
  return sr
end

function SRow(a::NfRelElem)
  sr = SRow(base_field(parent(a)))
  for i=0:length(a.data)
    c = coeff(a.data, i)
    if !iszero(c)
      push!(sr.pos, i+1)
      push!(sr.values, c)
    end
  end
  return sr
end

function minpoly_dense(a::NfRel_nsElem)
  K = parent(a)
  n = degree(K)
  k = base_field(K)
  M = zero_matrix(k, degree(K)+1, degree(K))
  z = one(K)
  elem_to_mat_row!(M, 1, z)
  z *= a
  elem_to_mat_row!(M, 2, z)
  i = 2
  while true
    if n % (i-1) == 0 && rank(M) < i
      N = nullspace(sub(M, 1:i, 1:ncols(M))')
      @assert N[1] == 1
      f = PolynomialRing(k,"t", cached=false)[1]([N[2][j, 1] for j=1:i])
      return f*inv(lead(f))
    end
    z *= a
    elem_to_mat_row!(M, i+1, z)
    i += 1
  end
end

function Base.Matrix(a::SMat)
  A = zero_matrix(base_ring(a), nrows(a), ncols(a))
  for i = 1:nrows(a)
    for (k, c) = a[i]
      A[i, k] = c
    end
  end
  return A
end

function minpoly_sparse(a::NfRel_nsElem)
  K = parent(a)
  n = degree(K)
  k = base_field(K)
  M = sparse_matrix(k)
  z = one(K)
  sz = SRow(z)
  i = 0
  push!(sz.values, k(1))
  push!(sz.pos, n+i+1)
  push!(M, sz)
  z *= a
  sz = SRow(z)
  i = 1
  kt, t = PolynomialRing(k, "t", cached = false)
  f = kt()
  while true
    if n % i == 0
      echelon!(M)
      fl, so = can_solve_ut(sub(M, 1:i, 1:n), sz)
      if fl
        so = mul(so, sub(M, 1:i, n+1:ncols(M)))
        # TH: If so is the zero vector, we cannot use the iteration,
        # so we do it by hand.
        if length(so.pos) == 0
          f = t^i
        else
          f = t^i - sum(c*t^(k-1) for (k,c) = so)
        end
        return f
      end
    end
    push!(sz.values, k(1))
    push!(sz.pos, n+i+1)
    push!(M, sz)
    z *= a
    sz = SRow(z)
    i += 1
    if i > degree(parent(a))
      error("too large")
    end
  end
end

function minpoly(a::NfRel_nsElem)
  return minpoly_sparse(a)
end

function minpoly(a::NfRel_nsElem{nf_elem}, ::FlintRationalField)
  f = minpoly(a)
  n = norm(f)
  g = gcd(n, derivative(n))
  if isone(g)
    return n
  end
  n = divexact(n, g)
  return n
end

absolute_minpoly(a::NfRel_nsElem{nf_elem}) = minpoly(a, FlintQQ)

function inv(a::NfRel_nsElem)
  if iszero(a)
    error("division by zero")
  end
  K = parent(a)
  f = minpoly(a)
  z = K(coeff(f, degree(f)))
  for i=degree(f)-1:-1:1
    z = z*a + coeff(f, i)
  end
  return -z*inv(coeff(f, 0))
end

function charpoly(a::NfRel_nsElem)
  f = minpoly(a)
  return f^div(degree(parent(a)), degree(f))
end

function norm(a::NfRel_nsElem)
  f = minpoly(a)
  return (-1)^degree(parent(a)) * coeff(f, 0)^div(degree(parent(a)), degree(f))
end

function tr(a::NfRel_nsElem)
  f = minpoly(a)
  return -coeff(f, degree(f)-1)*div(degree(parent(a)), degree(f))
end

#TODO: also provide a sparse version
function representation_matrix(a::NfRel_nsElem)
  K = parent(a)
  b = basis(K)
  k = base_field(K)
  M = zero_matrix(k, degree(K), degree(K))
  for i=1:degree(K)
    elem_to_mat_row!(M, i, a*b[i])
  end
  return M
end

@inline ngens(K::NfRel_ns) = length(K.pol)

mutable struct NfRelToNfRel_nsMor{T} <: Map{NfRel{T}, NfRel_ns{T}, HeckeMap, NfRelToNfRel_nsMor}
  header::MapHeader{NfRel{T}, NfRel_ns{T}}
  prim_img::NfRel_nsElem{T}
  emb::Array{NfRelElem{T}, 1}
  coeff_aut::NfToNfMor

  function NfRelToNfRel_nsMor(K::NfRel{T}, L::NfRel_ns{T}, a::NfRel_nsElem{T}, emb::Array{NfRelElem{T}, 1}) where {T}
    function image(x::NfRelElem{T})
      # x is an element of K
      f = data(x)
      # First evaluate the coefficients of f at a to get a polynomial over L
      # Then evaluate at b
      return f(a)
    end

    function preimage(x::NfRel_nsElem{T})
      return msubst(x.data, emb)
    end

    z = new{T}()
    z.prim_img = a
    z.emb = emb
    z.header = MapHeader(K, L, image, preimage)
    return z
  end

  function NfRelToNfRel_nsMor(K::NfRel{T}, L::NfRel_ns{T}, a::NfRel_nsElem{T}) where {T}
    function image(x::NfRelElem{T})
      # x is an element of K
      f = data(x)
      # First evaluate the coefficients of f at a to get a polynomial over L
      # Then evaluate at b
      return f(a)
    end

    z = new{T}()
    z.prim_img = a
    z.header = MapHeader(K, L, image)
    return z
  end

  function NfRelToNfRel_nsMor(K::NfRel{T}, L::NfRel_ns{T}, aut::Map, a::NfRel_nsElem{T}) where {T}
    aut = NfToNfMor(domain(aut), codomain(aut), aut(gen(domain(aut))))
    function image(x::NfRelElem{T})
      # x is an element of K
      f = deepcopy(data(x))
      for i=0:degree(f)
        setcoeff!(f, i, aut(coeff(f, i)))
      end
      # First evaluate the coefficients of f at a to get a polynomial over L
      # Then evaluate at b
      return f(a)
    end

    z = new{T}()
    z.prim_img = a
    z.coeff_aut = aut
    z.header = MapHeader(K, L, image)
    return z
  end

end

mutable struct NfRel_nsToNfRel_nsMor{T} <: Map{NfRel_ns{T}, NfRel_ns{T}, HeckeMap, NfRel_nsToNfRel_nsMor}
  header::MapHeader{NfRel_ns{T}, NfRel_ns{T}}
  emb::Array{NfRel_nsElem{T}, 1}
  coeff_aut::NfToNfMor

  function NfRel_nsToNfRel_nsMor(K::NfRel_ns{T}, L::NfRel_ns{T}, emb::Array{NfRel_nsElem{T}, 1}) where {T}
    function image(x::NfRel_nsElem{T})
      # x is an element of K
      # First evaluate the coefficients of f at a to get a polynomial over L
      # Then evaluate at b
      return msubst(x.data, emb)
    end

    z = new{T}()
    z.coeff_aut = id_hom(base_field(K))
    z.emb = emb
    z.header = MapHeader(K, L, image)
    return z
  end

  function NfRel_nsToNfRel_nsMor(K::NfRel_ns{T}, L::NfRel_ns{T}, aut::Map, emb::Array{NfRel_nsElem{T}, 1}) where {T}
    function image(x::NfRel_nsElem{T})
      # x is an element of K
      # First evaluate the coefficients of f at a to get a polynomial over L
      # Then evaluate at b
      f = x.data
      Kbxyz = parent(f)
      k = nvars(Kbxyz)
      Lbxyz = PolynomialRing(base_field(L), k, cached = false)[1]
      coeffs = Vector{T}(undef, length(f.coeffs))
      for i = 1:length(coeffs)
        coeffs[i] = aut(f.coeffs[i])
      end
      g = Lbxyz(coeffs, f.exps)
      return msubst(g, emb)
    end

    z = new{T}()
    z.emb = emb
    z.coeff_aut = aut
    z.header = MapHeader(K, L, image)
    return z
  end
end

id_hom(K::NfRel_ns) = NfRel_nsToNfRel_nsMor(K, K, gens(K))

function Base.hash(f::NfRel_nsToNfRel_nsMor{T}, u::UInt64) where T
  #I combine the hash functions for the automorphism of the base field and the hash function for the images of the generators.
  a = hash(f.coeff_aut, u)
  for i = 1:length(f.emb)
    a = hash(f.emb[i], a)
  end
  return a
end

function Base.:(*)(f::NfRel_nsToNfRel_nsMor{T}, g::NfRel_nsToNfRel_nsMor{T}) where {T}
  codomain(f) == domain(g) || throw("Maps not compatible")

  a = gens(domain(g))
  return NfRel_nsToNfRel_nsMor(domain(g), codomain(f), f.coeff_aut * g.coeff_aut, [ g(f(x)) for x in a])
end

function Base.:(==)(f::NfRel_nsToNfRel_nsMor{T}, g::NfRel_nsToNfRel_nsMor{T}) where {T}
  if domain(f) != domain(g) || codomain(f) != codomain(g)
    return false
  end

  L = domain(f)
  K = base_field(L)

  if f.coeff_aut.prim_img != g.coeff_aut.prim_img
    return false
  end

  for i = 1:ngens(L)
    if f.emb[i] != g.emb[i]
      return false
    end
  end

  return true
end

@inline ngens(R::Nemo.Generic.MPolyRing) = R.num_vars

function _prod(A, b)
  for a in A
    b = b * a
  end
  return b
end

#aparently, should be called evaluate, talk to Bill...
#definitely non-optimal, in particular for automorphisms
function msubst(f::Generic.MPoly{T}, v::Array{NfRelElem{T}, 1}) where T
  k = base_ring(parent(f))
  n = length(v)
  @assert n == ngens(parent(f))
  r = zero(parent(v[1]))
  L = parent(v[1])
  for i=1:length(f)
    #@show prod(v[j]^f.exps[j, i] for j=1:n)
    s = _prod((v[j]^f.exps[n - j + 1, i] for j=1:n), one(L))
    r += f.coeffs[i]* s
  end
  return r
end
function msubst(f::Generic.MPoly{T}, v::Array{NfRel_nsElem{T}, 1}) where T
  k = base_ring(parent(f))
  n = length(v)
  @assert n == ngens(parent(f))
  r = zero(k)
  for i=1:length(f)
    r += f.coeffs[i]*prod(v[j]^f.exps[n - j + 1, i] for j=1:n)
  end
  return r
end

function simple_extension(K::NfRel_ns{T}) where {T}
  n = ngens(K)
  g = gens(K)
  if n == 1
    fl, p = isunivariate(K.pol[1])
    Ks, gKs = number_field(p, cached = false, check = false)
    return Ks, NfRelToNfRel_nsMor(Ks, K, g[1], [gKs])
  end
  if lcm([total_degree(K.pol[i]) for i = 1:length(K.pol)]) == degree(K)
    #The sum of the primitive elements is the right element
    pe = sum(g[i] for i = 1:length(g))
    f = minpoly(pe)
  else
    pe = g[1]
    i = 1
    f = minpoly(pe)
    #todo: use resultants rather than minpoly??
    while i < n
      i += 1
      pe += g[i]
      f = minpoly(pe)
      while degree(f) < prod(total_degree(K.pol[k]) for k=1:i)
        pe += g[i]
        f = minpoly(pe)
      end
    end
  end
  Ka, a = number_field(f, check = false)
  k = base_field(K)
  M = zero_matrix(k, degree(K), degree(K))
  z = one(K)
  elem_to_mat_row!(M, 1, z)
  elem_to_mat_row!(M, 2, pe)
  mul!(z, z, pe)
  for i=3:degree(K)
    mul!(z, z, pe)
    elem_to_mat_row!(M, i, z)
  end
  N = zero_matrix(k, 1, degree(K))
  b1 = basis(Ka)
  emb = Array{NfRelElem{T}, 1}(undef, n)
  for i = 1:n
    elem_to_mat_row!(N, 1, g[i])
    s = solve(M', N')
    emb[i] = zero(Ka)
    for j = 1:degree(Ka)
      emb[i] += b1[j]*s[j, 1]
    end
  end

  return Ka, NfRelToNfRel_nsMor(Ka, K, pe, emb)
end

@doc Markdown.doc"""
    simple_extension(K::NfRel_ns{nf_elem}, FlintQQ) -> AnticNumberField, Map, Map

    absolute_field(K::NfRel_ns{nf_elem}) -> AnticNumberField, Map, Map

Compute an isomorphic field as an extension of $Q$ together with the isomorphism
(1st map) and the embedding of the base field (2nd map).
"""
function simple_extension(K::NfRel_ns{nf_elem}, ::FlintRationalField)
  Ks, mp = simple_extension(K)
  Ka, m1, m2 = absolute_field(Ks)
  return Ka, m1*mp, m2
end

absolute_field(K::NfRel_ns{nf_elem}) = simple_extension(K, FlintQQ)

#trivia, missing in NfRel
function basis(K::NfRel)
  a = gen(K)
  z = one(K)
  b = [z, a]
  while length(b) < degree(K)
    push!(b, b[end]*a)
  end
  return b
end

function Base.one(a::NfRelElem)
  return one(parent(a))
end

function Base.copy(a::NfRelElem)
  return parent(a)(a.data)
end

function Nemo.discriminant(K::NfRel_ns)
  p = [isunivariate(x)[2] for x = K.pol]
  d = discriminant(p[1])
  n = degree(p[1])
  for i=2:length(p)
    d = d^degree(p[i]) * discriminant(p[i])^n
    n *= degree(p[i])
  end
  return d
end

function Nemo.discriminant(K::NfRel_ns, ::FlintRationalField)
  d = norm(discriminant(K)) * discriminant(base_field(K))^degree(K)
  return d
end


text:
#!/usr/bin/env julia

using MPI
using PencilArrays
using Random
using Test

import ArrayInterface:
    ArrayInterface,
    StaticInt,
    StaticBool,
    contiguous_axis,
    contiguous_axis_indicator,
    contiguous_batch_size,
    dense_dims,
    stride_rank

struct DummyArray{T,N} <: AbstractArray{T,N}
    dims :: Dims{N}
    DummyArray{T}(::UndefInitializer, dims) where {T} = new{T,length(dims)}(dims)
end

Base.size(x::DummyArray) = x.dims
Base.getindex(::DummyArray{T}, ind...) where {T} = zero(T)
Base.strides(x::DummyArray) = Base.size_to_strides(1, size(x)...)

function non_dense_array(::Type{T}, dims) where {T}
    # Only the first dimension is dense: (True, False, False, ...).
    N = length(dims)
    dims_parent = ntuple(d -> (d - 1) + dims[d], Val(N))
    up = view(Array{T}(undef, dims_parent), Base.OneTo.(dims)...)
    @assert dense_dims(up) === ntuple(d -> StaticBool(d == 1), Val(ndims(up)))
    up
end

function non_contiguous_array(::Type{T}, dims) where {T}
    N = length(dims)
    dims_parent = (2, dims...)  # we take the slice [1, :, :, ...]
    up = view(Array{T}(undef, dims_parent), 1, ntuple(d -> Colon(), Val(N))...)
    @assert contiguous_axis(up) === StaticInt(-1)
    @assert size(up) == dims
    @assert ArrayInterface.size(up) == dims
    up
end

function test_array_interface(p::Pencil)
    # Test different kinds of parent arrays
    dims_mem = size_local(p, MemoryOrder())
    up_regular = Array{Float64}(undef, dims_mem)
    up_noncontig = non_contiguous_array(Float64, dims_mem)
    up_nondense = non_dense_array(Float64, dims_mem)
    up_dummy = DummyArray{Float64}(undef, dims_mem)

    parents = (
        up_regular,
        up_noncontig,
        up_nondense,
        up_dummy,
    )

    @testset "Parent $(typeof(up))" for up in parents
        u = PencilArray(p, up)

        @test ArrayInterface.parent_type(u) === typeof(up)
        @test ArrayInterface.known_length(u) === missing
        @test !ArrayInterface.can_change_size(u)
        @test ArrayInterface.ismutable(u)
        @test ArrayInterface.can_setindex(u)
        @test ArrayInterface.aos_to_soa(u) === u
        @test ArrayInterface.fast_scalar_indexing(u)
        @test !ArrayInterface.isstructured(u)

        # Compare outputs with equivalent PermutedDimsArray
        iperm = inv(permutation(u))
        vp = PermutedDimsArray(up, iperm)

        functions = (
            contiguous_axis, contiguous_axis_indicator,
            contiguous_batch_size, stride_rank, dense_dims,
            ArrayInterface.size, ArrayInterface.strides, ArrayInterface.offsets,
        )

        for f in functions
            @inferred f(u)
            @test f(u) === f(vp)
        end
    end

    nothing
end

MPI.Init()
Nxyz = (11, 21, 32)
comm = MPI.COMM_WORLD
Nproc = MPI.Comm_size(comm)
myrank = MPI.Comm_rank(comm)

let dev_null = @static Sys.iswindows() ? "nul" : "/dev/null"
    MPI.Comm_rank(comm) == 0 || redirect_stdout(open(dev_null, "w"))
end

rng = MersenneTwister(42 + myrank)

topo = MPITopology(comm, Val(2))

pen1 = Pencil(topo, Nxyz, (2, 3))
pen2 = Pencil(pen1, decomp_dims=(1, 3), permute=Permutation(2, 1, 3))
pen3 = Pencil(pen2, decomp_dims=(1, 2), permute=Permutation(3, 2, 1))
pens = (pen1, pen2, pen3)

@testset "ArrayInterface -- Pencil$(decomposition(p))" for p in pens
    test_array_interface(p)
end


text:
struct HGrid{T <: Real, N  <: Integer}
                   nxh :: N                                     # Number of grid cells in x-direction in HGrid
                   nyh :: N                                     # Number of grid cells in y-direction in HGrid
                 mask :: Array{Bool,2}                         # Mask specifying the model domain
                    n :: N                                     # Total number of cells in the model domain
                 crop :: Diagonal{T,Array{T,1}}                # Crop matrix: diagonal matrix with mask entries on diag
                 samp :: SparseMatrixCSC{T,N}                  # Sampling matrix: take full domain to model domain
               spread :: SparseMatrixCSC{T,N}                  # Sparse form of the sampling matrix
                    b :: Array{T,2}                            # Bed elevation
                    h :: Array{T,2}                            # Ice thickness
                    s :: Array{T,2}                            # Current surface elevation
                 dhdt :: Array{T,2}                            # Time rate of change of thickness
         accumulation :: Array{T,2}                            # Accumulation rate
           basal_melt :: Array{T,2}                            # Basal melt rate
                  haf :: Array{T,2}                            # Grid cell height above floatation
    grounded_fraction :: Array{T,2}                            # Grid cell grounded fraction
                 dsdh :: Array{T,2}                            # Cange of surface elevation per unit thickness change
    shelf_strain_rate :: Array{T,2}                            # Strain rate appropriate for shelf (no basal drag)
             av_speed :: Array{T,2}                            # Depth averaged speed
                    u :: Array{T,2}                            # Depth averaged x-velocity
                    v :: Array{T,2}                            # Depth averaged y-velocity
                   us :: Array{T,2}                            # x-velocity at the surface
                   vs :: Array{T,2}                            # y-velocity at the surface
                   ub :: Array{T,2}                            # x-velocity at the bed
                   vb :: Array{T,2}                            # y-velocity at the bed
            bed_speed :: Array{T,2}                            # Ice speed at the bed
           weertman_c :: Array{T,2}                            # Weertman drag coefficients
                    β :: Array{T,2}                            # Raw β value (eqn 8 in Arthern 2015 JGeophysRes)
                 βeff :: Array{T,2}                            # Effective β value (eqn 12 in Arthern 2015 JGeophysRes)
                 τbed :: Array{T,2}                            # Stress at the bed
                  ηav :: Array{T,2}                            # Depth averaged viscosity
              quad_f1 :: Array{T,2}                            # F1 quadratrue field (eqn 7 in Arthern 2015 JGeophysRes)
              quad_f2 :: Array{T,2}                            # F2 quadrature field (eqn 7 in Arthern 2015 JGeophysRes)
             dneghηav :: Base.RefValue{Diagonal{T,Array{T,1}}} # Rheological operator (-h × ηav)
            dimplicit :: Base.RefValue{Diagonal{T,Array{T,1}}} # Rheological operator (-ρi × g × dt × dshs)
end

"""
    HGrid(;
            nxh,
            nyh,
            mask = trues(nxh,nyh),
            b,
            h,
            ηav = zeros(nxh,nyh))

Construct a WAVI.jl HGrid with size (nxh,nyh)
HGrid stores fields that are defined on the problem's H grid.
(Co-ordinates of HGrid stored in a Grid under xxh, yyh fields)

Keyword arguments
=================
    - 'nxh': (required) Number of grid cells in x-direction in HGrid (should be same as grid.nx)
            Note that we store the grid size here, even though it can be easily inferred from grid, to increase transparency in velocity solve.
    - 'nyh': (required) Number of grid cells in y-direction in HGrid (should be same as grid.ny)
    - 'mask': Mask specifying the model domain
    - 'b': (requried) Bed elevation (bottom bathymetry)
    - 'h': (required) initial thickness of the ice
    - 'ηav': depth averaged visosity initially
"""

function HGrid(;
                nxh,
                nyh,
                mask = trues(nxh,nyh),
                b,
                h = zeros(nxh,nyh),
                ηav = zeros(nxh,nyh))

    #check the sizes of inputs
    (size(mask) == size(b) == size(h) == size(ηav) == (nxh,nyh)) || throw(DimensionMismatch("Sizes of inputs to UGrid must all be equal to nxh x nyh (i.e. $nxh x $nyh)"))

    #construct operators
    n = count(mask);
    crop = Diagonal(float(mask[:]));
    samp = crop[mask[:],:];
    spread = sparse(samp');
    dneghηav = Ref(crop*Diagonal(zeros(nxh*nyh))*crop)
    dimplicit = Ref(crop*Diagonal(zeros(nxh*nyh))*crop)

    #construct quantities not passed
    s = zeros(nxh,nyh)
    dhdt = zeros(nxh,nyh)
    accumulation = zeros(nxh,nyh)
    basal_melt = zeros(nxh,nyh)
    haf = zeros(nxh,nyh)
    grounded_fraction = ones(nxh,nyh)
    dsdh = ones(nxh,nyh)
    shelf_strain_rate = zeros(nxh,nyh)
    av_speed = zeros(nxh,nyh)
    u = zeros(nxh,nyh)
    v = zeros(nxh,nyh)
    us = zeros(nxh,nyh)
    vs = zeros(nxh,nyh)
    ub = zeros(nxh,nyh)
    vb= zeros(nxh,nyh)
    bed_speed = zeros(nxh,nyh)
    weertman_c = zeros(nxh,nyh)
    β = zeros(nxh,nyh)
    βeff = zeros(nxh,nyh)
    τbed = zeros(nxh,nyh)
    quad_f1 = zeros(nxh,nyh)
    quad_f2 = h./(3*ηav)

    #check sizes of everything
    @assert size(mask)==(nxh,nyh); @assert mask == clip(mask)
    @assert n == count(mask)
    @assert crop == Diagonal(float(mask[:]))
    @assert samp == crop[mask[:],:]
    @assert spread == sparse(samp')
    @assert size(b)==(nxh,nyh)
    @assert size(h)==(nxh,nyh)
    @assert size(s)==(nxh,nyh)
    @assert size(dhdt)==(nxh,nyh)
    @assert size(accumulation)==(nxh,nyh)
    @assert size(basal_melt)==(nxh,nyh)
    @assert size(haf)==(nxh,nyh)
    @assert size(grounded_fraction)==(nxh,nyh)
    @assert size(dsdh)==(nxh,nyh)
    @assert size(shelf_strain_rate)==(nxh,nyh)
    @assert size(u)==(nxh,nyh)
    @assert size(v)==(nxh,nyh)
    @assert size(av_speed)==(nxh,nyh)
    @assert size(ub)==(nxh,nyh)
    @assert size(vb)==(nxh,nyh)
    @assert size(us)==(nxh,nyh)
    @assert size(vs)==(nxh,nyh)
    @assert size(bed_speed)==(nxh,nyh)
    @assert size(weertman_c)==(nxh,nyh)
    @assert size(β)==(nxh,nyh)
    @assert size(βeff)==(nxh,nyh)
    @assert size(τbed)==(nxh,nyh)
    @assert size(quad_f1)==(nxh,nyh)
    @assert size(quad_f2)==(nxh,nyh)
    @assert size(ηav)==(nxh,nyh)

    #make sure boolean type rather than bitarray
    mask = convert(Array{Bool,2}, mask)

return HGrid(
            nxh,
            nyh,
            mask,
            n,
            crop,
            samp,
            spread,
            b,
            h,
            s,
            dhdt,
            accumulation,
            basal_melt,
            haf,
            grounded_fraction,
            dsdh,
            shelf_strain_rate,
            av_speed,
            u,
            v,
            us,
            vs,
            ub,
            vb,
            bed_speed,
            weertman_c,
            β,
            βeff,
            τbed,
            ηav,
            quad_f1,
            quad_f2,
            dneghηav,
            dimplicit)
end


text:
module tree
import Base: @kwdef
import MyterialColors: yellow, orange, red, blue
using InteractiveUtils

import Term: loop_last,
        escape_brackets,
        fillin,
        highlight,
        int,
        theme,
        textlen,
        truncate
import ..segment: Segment
import ..measure: Measure
import ..renderables: AbstractRenderable
import ..style: apply_style
import ..layout: vstack, pad, hLine
import ..panel: Panel

export Tree

# ---------------------------------------------------------------------------- #
#                                    GUIDES                                    #
# ---------------------------------------------------------------------------- #

treeguides = Dict(
    :standardtree =>("    ", "│   ", "├── ", "└── "),
    :boldtree     =>("    ", "┃   ", "┣━━ ", "┗━━ "),
    :asciitree    =>("    ", "|   ", "+-- ", "`-- "),
)

"""
    TreeGuides

Store strings to make up a `Tree`'s guides (the line elements showing connections).
"""
struct TreeGuides
    space::String
    vline::String
    branch::String
    leaf::String
end

"""
    TreeGuides(guides_type::Symbol, style::String)

Get tree guides with style information applied
"""
TreeGuides(guides_type::Symbol, style::String) = TreeGuides(
    map(
        (g)->apply_style("[$style]$g[/$style]"),
        treeguides[guides_type]
    )...
)

# ---------------------------------------------------------------------------- #
#                                     TREE                                     #
# ---------------------------------------------------------------------------- #

# ----------------------------------- leaf ----------------------------------- #

"""
    asleaf

Style an object to render it as a a string
"""
function asleaf end

asleaf(x) = truncate(highlight(x), 22)
asleaf(x::Nothing) = nothing
asleaf(x::AbstractVector) = truncate((escape_brackets ∘ string)(x), 22)
asleaf(x::AbstractString) = truncate(highlight(x, :string), 22)

"""
    Leaf

End items in a `Tree`. No sub-trees.
"""
struct Leaf
    name::Union{Nothing, String}
    text::Union{Nothing, String}
end

# ----------------------------------- tree ----------------------------------- #
"""
    Tree

A tree is composed of nodes (other trees) and leaves (end items).
It renders as a hierarchical structure with lines (guides) connecting the various elements
"""
@kwdef struct Tree <: AbstractRenderable
    segments::Union{Nothing, Vector{Segment}} = nothing
    measure::Union{Nothing, Measure} = nothing

    name::String
    level::Int
    nodes::Vector{Tree}
    leaves::Vector{Leaf}

    title_style::String = theme.tree_title_style
    node_style::String = theme.tree_node_style
    leaf_style::String = theme.tree_leaf_style
    guides_style::String = theme.tree_guide_style
    guides_type::Symbol = :standardtree
end

"""
Show/render a `Tree`
"""
function Base.show(io::IO, tree::Tree)
    if io != stdout
        print(io, "Tree: $(length(tree.nodes)) nodes, $(length(tree.leaves)) leaves")
    else
        for seg in tree.segments
            println(io, seg)
        end
    end
end

"""
    Tree(data::Union{Dict, Pair}; level=0, title::String="tree", kwargs...)

Construct a `Tree` out of a `Dict`. Recursively handle nested `Dict`s.
"""
function Tree(
        data::Union{Dict, Pair};
        level=0,
        title::String="tree",
        kwargs...
    )

    # initialize
    nodes::Vector{Tree} = []
    leaves::Vector{Leaf} = []

    # go over all entries
    for (k, v) in zip(keys(data), values(data))
        if v isa Dict
            push!(nodes, Tree(v; level=level+1, title=truncate(string(k), 22)))
        elseif v isa Pair
            k = isnothing(v.first) ? nothing : truncate(v.first, 22)
            push!(leaves, Leaf(v, asleaf(v.second)))
        else
            k = isnothing(k) ? nothing : truncate(string(k), 22)
            push!(leaves, Leaf(k, asleaf(v)))
        end
    end

    # if we're handling the first tree, render it. Otherwise parse nested trees.
    if level > 0
        # we don't need to render
        return Tree(; name=title, level=level, nodes=nodes, leaves=leaves, kwargs...)
    else
        # render and get measure
        segments = render(
                Tree(;
                    name=title,
                    level=level,
                    nodes=nodes,
                    leaves=leaves,
                    kwargs...
            )
        )
        measure = Measure(segments)

        return Tree(;
            segments=segments,
            measure=measure,
            name=truncate(title, 22),
            level=level,
            nodes=nodes,
            leaves=leaves,
            kwargs...
        )
    end
end

# ---------------------------------- render ---------------------------------- #

"""
    render(tree::Tree)::Vector{Segment}

Render a `Tree` into segments. Recursively handle nested trees.

"""
function render(tree::Tree; prevguides::String="", lasttree=false, waslast=[], guides=nothing)::Vector{Segment}
    guides = isnothing(guides) ? TreeGuides(tree.guides_type, tree.guides_style) : guides
    hasleaves = length(tree.leaves) > 0

    segments::Vector{Segment}=[]

    """
        Add a segment to the segments vector
    """
    function _add(x::String, style)
        push!(segments, Segment(x, style))
    end
    _add(x::String) = _add(x, "default")

    # ------------------------------ render in parts ----------------------------- #
    # render initial part
    if tree.level == 0
        header_text = "[$(tree.title_style)]$(tree.name)[/$(tree.title_style)]"
        header = (" " * header_text * " ") / hLine(textlen(tree.name)+2; style="$(tree.title_style) dim", box=:HEAVY)

        append!(segments, header.segments)
    else
        _pre_guides = ""
        for (n, (l, last)) in enumerate(loop_last(waslast))
            # ugly, get the correct sequence of guides based on where we are/what came before
            if last
                _end = lasttree ? guides.leaf : guides.branch
                _pre_guides *= l ? _end : guides.space

            else
                _end = n == length(waslast) ? guides.branch : guides.vline
                _pre_guides *= lasttree ? (l ? guides.leaf : guides.vline) : _end
            end
        end

        _add(_pre_guides * "[$(tree.node_style)]$(tree.name)[/$(tree.node_style)]")
    end
    tree.level == 0 && _add(prevguides * guides.vline)

    # render sub-trees
    for (last, node) in loop_last(tree.nodes)
        # check if it's the last entry in the tree
        lasttree = last && !hasleaves

        # get the appropriate guides
        if lasttree
            prev = prevguides * guides.space
        else
            prev = prevguides * guides.vline
        end

        append!(segments, render(
            node;
            prevguides=prev,
            lasttree=lasttree,
            waslast=vcat(waslast, lasttree),
            guides=guides
            )
        )
        hasleaves && length(node.leaves) > 0 && _add(prevguides * guides.vline)
    end

    # render leaves
    if hasleaves
        for (last, leaf) in loop_last(tree.leaves)
            seg = last ? guides.leaf : guides.branch
            if isnothing(leaf.text)
                k = isnothing(leaf.name) ? "" : highlight(leaf.name)
                v = ""
            else
                k = isnothing(leaf.name) ? "" : "[$(tree.leaf_style)]$(leaf.name)[/$(tree.leaf_style)]: "
                v = leaf.text
            end
            _add(prevguides * seg * k * v)
        end
    end

    # left pad segments
    if tree.level == 0
        header_length = length(header.segments)
        padded_segments = vcat(
            header.segments...,
            pad(segments[header_length+1:end], int(header.measure.w/2 - 1))...
        )
        return fillin(padded_segments)
    else
        return segments
    end

end

# ---------------------------------------------------------------------------- #
#                                HIERARCHY TREE                                #
# ---------------------------------------------------------------------------- #
"""
Apply style for the type whose hierarchy Tree we are making
"""
style_T(T) = "[orange1 italic underline]$T[/orange1 italic underline]"

"""
    make_hierarchy_dict(x::Vector{DataType}, T::DataType, Tsubs::Dict)::Dict

Recursively create a dictionary with the types hierarchy for `T`.
`Tsubs` carries information about T's subtypes.
The Dict is made backwards. From  the deepest levels up.
"""
function make_hierarchy_dict(x::NTuple, T::DataType, Tsubs::Dict)::Dict
    data = Dict()
    prev = ""
    for (n, y) in enumerate(x)
        if n == 1
            continue
        elseif n < length(x)
            subs = Dict()
            for s in subtypes(y)
                if s == T
                    subs[style_T(s)] = Tsubs
                else
                    subs[string(s)] = nothing
                end
            end

            if n == 2
                data = subs
            else
                subs[prev] = data
                data = subs
            end

            prev = string(y)
        end
    end
    return data
end

"""
    Tree(T::DataType)::Tree

Construct a `Tree` visualization of `T`'s types hierarchy
"""
function Tree(T::DataType)::Tree
    # create a dictionary of types hierarchy
    subs = Dict(string(s)=>nothing for s in subtypes(T))
    data = make_hierarchy_dict(supertypes(T), T, subs)

    return Tree(
            data;
            # title=string(supertypes(T)[end]),
            title="Any",
            title_style="bright_green italic",
            guides_style="green dim",
        )
end

end