text:
using IICGP
using Cambrian
using CartesianGeneticProgramming
using Test

import Cambrian.mutate # function extension

function test_fitness_matrix(e::DualCGPGAEvo)
    esize = (length(e.encoder_sympop), length(e.controller_sympop))
    @test size(e.fitness_matrix) == esize
    nrows, ncols = size(e.fitness_matrix)
    # 1. Test at least 1 eval per row col
    for r in eachrow(e.fitness_matrix)
        @test maximum(r) > -Inf
    end
    for c in eachcol(e.fitness_matrix)
        @test maximum(c) > -Inf
    end
    # 2. Test n_eval eval
    @test length(filter(x->x>-Inf, e.fitness_matrix)) == e.n_eval
    # 3. Test n_elites elites are selected
    @test sum(e.elites_matrix) == e.n_elite
    # 4. Test elites are evaluated and have higher scores
    for i in eachindex(e.elites_matrix)
        if e.elites_matrix[i]
            @test e.fitness_matrix[i] > -Inf
            for j in eachindex(e.fitness_matrix)
                if !e.elites_matrix[j]
                    @test e.fitness_matrix[i] >= e.fitness_matrix[j]
                end
            end
        end
    end
    # 5. Test all pairs in eval_matrix are evaluated
    for i in eachindex(e.elites_matrix)
        if e.elites_matrix[i]
            @test e.fitness_matrix[i] > -Inf
        end
    end
end

function test_ind_fitnesses(e::DualCGPGAEvo)
    # Test that individual's fitness is set according to the fitness matrix
    for i in eachindex(e.encoder_sympop)
        @test length(e.encoder_sympop[i].fitness) == e.config.d_fitness
        @test e.encoder_sympop[i].fitness[1] > -Inf
        @test e.encoder_sympop[i].fitness[1] == maximum(e.fitness_matrix[i,:])
    end
    for j in eachindex(e.controller_sympop)
        @test length(e.controller_sympop[j].fitness) == e.config.d_fitness
        @test e.controller_sympop[j].fitness[1] > -Inf
        @test e.controller_sympop[j].fitness[1] == maximum(e.fitness_matrix[:,j])
    end
end

cfg_filename = string(@__DIR__, "/dualcgpga_test.yaml")
game = "gravitar"
mcfg, ecfg, ccfg, reducer, bootstrap = IICGP.dualcgp_config(cfg_filename, game)
logid = "2021-testlogid"
resdir = dirname(@__DIR__)
function mutate(ind::CGPInd, ind_type::String)
    if ind_type == "encoder"
        return goldman_mutate(ecfg, ind, init_function=IPCGPInd)
    elseif ind_type == "controller"
        return goldman_mutate(ccfg, ind)
    end
end
fit(e::CGPInd, c::CGPInd, seed::Int64) = [e.chromosome[1] * c.chromosome[1]]
n_iter = 4

@testset "CGP GA Evaluation" begin
    evo = IICGP.DualCGPGAEvo(mcfg, ecfg, ccfg, fit, logid, resdir)
    for iter in 1:n_iter
        evaluate(evo)
        test_fitness_matrix(evo)
        test_ind_fitnesses(evo)
        populate(evo)
    end
end


text:
import .ForwardDiff: Dual, value, partials, Partials # The dot in .ForwardDiff is an artefact of using Requires.jl

"""
    switch_representation(d::Dual{T, V, N}) where {T, V <: AbstractParticles, N}

Goes from Dual{Particles} to Particles{Dual}
"""
function switch_representation(d::Dual{T,V,N}) where {T,V<:AbstractParticles,N}
    part = partials(d)
    MonteCarloMeasurements.nakedtypeof(V)([Dual{T}(value(d).particles[i], ntuple(j->part[j].particles[i], N)) for i ∈ 1:nparticles(V)])
end

# function switch_representation(p::Particles{Dual{T,V,N},NP}) where {T,V<:AbstractParticles,N,NP}
#     Dual{T}(Particles(value.(p.particles)), Particles(partials.(p.particles)))
# end

const DualParticles = Dual{T,V,N} where {T,V<:AbstractParticles,N}
for ff in [maximum,minimum,std,var,cov,mean,median,quantile]
    f = nameof(ff)
    pname = Symbol("p"*string(f))
    m = Base.parentmodule(ff)
    @eval ($pname)(d::DualParticles) = ($pname)(switch_representation(d))
end

macro andreverse(ex)
    def = splitdef(ex)
    if haskey(def,:whereparams) && !isempty(def[:whereparams])
        quote
            $(esc(ex))
            $(esc(def[:name]))($(esc(def[:args][2])), $(esc(def[:args][1]))) where $(esc(def[:whereparams]...)) = $(esc(def[:body]))
        end
    else
        quote
            $(esc(ex))
            $(esc(def[:name]))($(esc(def[:args][2])), $(esc(def[:args][1]))) = $(esc(def[:body]))
        end
    end
end

for PT in (Particles, StaticParticles)
    @eval begin
        @andreverse function Base.:(*)(p::$PT, d::Dual{T}) where {T}
            Dual{T}($PT(p.particles .* value(d)), ntuple(i->$PT(p.particles .* partials(d)[i]) ,length(partials(d))))
        end

        @andreverse function Base.:(+)(p::$PT, d::Dual{T}) where {T}
            Dual{T}($PT(p.particles .+ value(d)), ntuple(i->$PT(0p.particles .+ partials(d)[i]) ,length(partials(d))))
        end

        function Base.:(-)(p::$PT, d::Dual{T}) where {T}
            Dual{T}($PT(p.particles .- value(d)), ntuple(i->$PT(0p.particles .- partials(d)[i]) ,length(partials(d))))
        end
        function Base.:(-)(d::Dual{T}, p::$PT) where {T}
            Dual{T}($PT(value(d) .- p.particles), ntuple(i->$PT(0p.particles .+ partials(d)[i]) ,length(partials(d))))
        end

        function Base.promote_rule(::Type{ForwardDiff.Dual{T,V,NP}}, ::Type{$PT{S, N}}) where {T, V, NP, S, N}
            VS = promote_type(V,S)
            Dual{T, $PT{VS, N}, NP}
            # Dual{T}($PT(fill(value(d), N)), ntuple(i->$PT(fill(partials(d)[i], N)) ,length(partials(d))))
        end

        # function Base.promote_rule(::Type{ForwardDiff.Dual{T, V, N}}, ::Type{$PT{T, N}}) where {T, V, N, T, N}
        #     Dual$PT{}
        # end
    end
end

# Base.hidigit(x::AbstractParticles, base) = Base.hidigit(mean(x), base) # To avoid stackoverflow in some printing situations
# Base.hidigit(x::Dual, base) = Base.hidigit(x.value, base) # To avoid stackoverflow in some printing situations
# Base.round(d::Dual, r::RoundingMode) = round(d.value,r)

ext: jl


text:
# ---
# title: 67. Add Binary
# id: problem67
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Math, String
# link: <https://leetcode.com/problems/add-binary/description/>
# hidden: true
# ---
#
# Given two binary strings `a` and `b`, return _their sum as a binary string_.
#
#
#
# **Example 1:**
#
#
#
#     Input: a = "11", b = "1"
#     Output: "100"
#
#
# **Example 2:**
#
#
#
#     Input: a = "1010", b = "1011"
#     Output: "10101"
#
#
#
#
# **Constraints:**
#
#   * `1 <= a.length, b.length <= 104`
#   * `a` and `b` consist only of `'0'` or `'1'` characters.
#   * Each string does not contain leading zeros except for the zero itself.
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/NodeNumbering.jl/blob/master/LICENSE

module NodeNumbering

include("calc_bw.jl")
include("create_adjacency_graph.jl")
include("node_degrees.jl")
include("RCM.jl")
include("renumbering.jl")
include("create_RCM_adjacency.jl")
include("adjacency_visualization.jl")

end


text:
# This file contains functions for the Profit Efficiency Additive DEA model
"""
    ProfitAdditiveDEAModel
An data structure representing a profit additive DEA model.
"""
struct ProfitAdditiveDEAModel <: AbstractProfitDEAModel
    n::Int64
    m::Int64
    s::Int64
    weights::Symbol
    monetary::Bool
    dmunames::Union{Vector{String},Nothing}
    eff::Vector
    lambda::SparseMatrixCSC{Float64, Int64}
    techeff::Vector
    alloceff::Vector
    normalization::Vector
    Xtarget::Matrix
    Ytarget::Matrix
end

"""
    deaprofitadd(X, Y, W, P, model)
Compute profit efficiency using data envelopment analysis weighted additive model for
inputs `X`, outputs `Y`, price of inputs `W`, and price of outputs `P`.

Model specification:
- `:Ones`: standard additive DEA model.
- `:MIP`: Measure of Inefficiency Proportions. (Charnes et al., 1987; Cooper et al., 1999)
- `:Normalized`: Normalized weighted additive DEA model. (Lovell and Pastor, 1995)
- `:RAM`: Range Adjusted Measure. (Cooper et al., 1999)
- `:BAM`: Bounded Adjusted Measure. (Cooper et al, 2011)
- `:Custom`: User supplied weights.

# Optional Arguments
- `rhoX`: matrix of weights of inputs. Only if `model=:Custom`.
- `rhoY`: matrix of weights of outputs. Only if `model=:Custom`.
- `names`: a vector of strings with the names of the decision making units.

# Examples
```jldoctest
julia> X = [2; 4; 8; 12; 6; 14; 14; 9.412];

julia> Y = [1; 5; 8; 9; 3; 7; 9; 2.353];

julia> W = [1; 1; 1; 1; 1; 1; 1; 1];

julia> P = [2; 2; 2; 2; 2; 2; 2; 2];

julia> deaprofitadd(X, Y, W, P, :Ones)
Profit Additive DEA Model
DMUs = 8; Inputs = 1; Outputs = 1
Weights = Ones; Returns to Scale = VRS
────────────────────────────────
   Profit  Technical  Allocative
────────────────────────────────
1   8.0      0.0        8.0
2   2.0      0.0        2.0
3   0.0      0.0        0.0
4   2.0      0.0        2.0
5   8.0      4.0        4.0
6   8.0      7.33333    0.666667
7   4.0      2.0        2.0
8  12.706    8.059      4.647
────────────────────────────────
```
"""
function deaprofitadd(X::Union{Matrix,Vector}, Y::Union{Matrix,Vector},
    W::Union{Matrix,Vector}, P::Union{Matrix,Vector},
    model::Symbol = :Default;
    rhoX::Union{Matrix,Vector,Nothing} = nothing, rhoY::Union{Matrix,Vector,Nothing} = nothing,
    monetary::Bool = false,
    names::Union{Vector{String},Nothing} = nothing,
    optimizer::Union{DEAOptimizer,Nothing} = nothing)::ProfitAdditiveDEAModel

    # Check parameters
    nx, m = size(X, 1), size(X, 2)
    ny, s = size(Y, 1), size(Y, 2)

    nw, mw = size(W, 1), size(W, 2)
    np, sp = size(P, 1), size(P, 2)

    if nx != ny
        throw(DimensionMismatch("number of rows in X and Y ($nx, $ny) are not equal"));
    end
    if nw != nx
        throw(DimensionMismatch("number of rows in W and X ($nw, $nx) are not equal"));
    end
    if np != ny
        throw(DimensionMismatch("number of rows in P and Y ($np, $ny) are not equal"));
    end
    if mw != m
        throw(DimensionMismatch("number of columns in W and X ($mw, $m) are not equal"));
    end
    if sp != s
        throw(DimensionMismatch("number of columns in P and Y ($sp, $s) are not equal"));
    end

    # Default behaviour
    if model == :Default
        # If no weights are specified use :Ones
        if rhoX === nothing && rhoY === nothing
            model = :Ones
        else
            model = :Custom
        end
    end

    # Get weights based on the selected model
    if model != :Custom
        # Display error if both model and weights are specified
        if rhoX !== nothing || rhoY !== nothing
            throw(ArgumentError("`rhoX` and `rhoY` not allowed if model != :Custom"));
        end

        # Get weights for selected model
        rhoX, rhoY = deaaddweights(X, Y, model, orient = :Graph)
    end

    if size(rhoX) != size(X)
        throw(DimensionMismatch("size of rhoX and X ($(size(rhoX)), $(size(X))) are not equal"));
    end
    if size(rhoY) != size(Y)
        throw(DimensionMismatch("size of rhoY and Y ($(size(rhoY)), $(size(Y))) are not equal"));
    end

    # Default optimizer
    if optimizer === nothing
        optimizer = DEAOptimizer(:LP)
    end

    # Get maximum profit targets and lambdas
    n = nx

    Xtarget, Ytarget, plambdaeff = deamaxprofit(X, Y, W, P, optimizer = optimizer)

    # Profit, technical and allocative efficiency
    maxprofit = sum(P .* Ytarget, dims = 2) .- sum(W .* Xtarget, dims = 2)
    pefficiency  = vec(maxprofit .- ( sum(P .* Y, dims = 2) .- sum(W .* X, dims = 2)))
    normalization = vec(minimum(hcat(W ./ rhoX, P ./ rhoY), dims = 2))
    techefficiency = efficiency(deaadd(X, Y, :Custom, rts = :VRS, rhoX = rhoX, rhoY = rhoY, optimizer = optimizer))

    if monetary
        techefficiency = techefficiency .* normalization
    else
        pefficiency = pefficiency ./ normalization
    end

    allocefficiency = pefficiency - techefficiency

    return ProfitAdditiveDEAModel(n, m, s, model, monetary, names, pefficiency, plambdaeff, techefficiency, allocefficiency, normalization, Xtarget, Ytarget)

end

function Base.show(io::IO, x::ProfitAdditiveDEAModel)
    compact = get(io, :compact, false)

    n = nobs(x)
    m = ninputs(x)
    s = noutputs(x)
    dmunames = names(x)

    eff = efficiency(x)
    techeff = efficiency(x, :Technical)
    alloceff = efficiency(x, :Allocative)

    if !compact
        print(io, "Profit Additive DEA Model \n")
        print(io, "DMUs = ", n)
        print(io, "; Inputs = ", m)
        print(io, "; Outputs = ", s)
        print(io, "\n")
        print(io, "Weights = ", string(x.weights))
        print(io, "; Returns to Scale = VRS")
        print(io, "\n")
        show(io, CoefTable(hcat(eff, techeff, alloceff), ["Profit", "Technical", "Allocative"], dmunames))
    end

end


text:
module StructuredQueries

using Compat

export  Cursor,
        Grouped,
        @with,
        source,
        graph

include("utils.jl")

# grouped
include("grouped/grouped.jl")
include("grouped/show.jl")

#verbs
include("verbs/verbs.jl")
include("verbs/primitives.jl")

include("verbs/expr/assignment_expr_ops.jl")
include("verbs/expr/scalar.jl")
include("verbs/expr/sym_analysis.jl")

include("verbs/process/generic.jl")
include("verbs/process/select.jl")
include("verbs/process/filter.jl")
include("verbs/process/orderby.jl")
include("verbs/process/groupby.jl")
include("verbs/process/join.jl")
include("verbs/process/summarize.jl")

#nodes
include("nodes/nodes.jl")
include("nodes/primitives.jl")
include("nodes/show.jl")

# cursors
include("cursor/cursor.jl")
include("cursor/primitives.jl")
include("cursor/show.jl")
include("cursor/macros.jl")
include("cursor/graph.jl")

# collect
include("collect/lift.jl")

end


text:
# Autogenerated wrapper script for FLANN_jll for x86_64-apple-darwin
export libflann, libflann_cpp

using CompilerSupportLibraries_jll
JLLWrappers.@generate_wrapper_header("FLANN")
JLLWrappers.@declare_library_product(libflann, "@rpath/libflann.1.9.dylib")
JLLWrappers.@declare_library_product(libflann_cpp, "@rpath/libflann_cpp.1.9.dylib")
function __init__()
    JLLWrappers.@generate_init_header(CompilerSupportLibraries_jll)
    JLLWrappers.@init_library_product(
        libflann,
        "lib/libflann.1.9.1.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@init_library_product(
        libflann_cpp,
        "lib/libflann_cpp.1.9.1.dylib",
        RTLD_LAZY | RTLD_DEEPBIND,
    )

    JLLWrappers.@generate_init_footer()
end  # __init__()


text:
#-----------------------------------------------------------------------------# Losses
l2regloss(y, yhat) = .5 * abs2(y - yhat)
deriv(::typeof(l2regloss), y, yhat) = yhat - y

l1regloss(y, yhat) = abs(y - yhat)
deriv(::typeof(l1regloss), y, yhat) = sign(yhat - y)

logisticloss(y, yhat) = log1p(exp(-y * yhat))
deriv(::typeof(logisticloss), y, yhat) = -y / (1 + exp(y * yhat))

l1hingeloss(y, yhat) = max(1 - y*yhat, zero(yhat))
deriv(::typeof(l1hingeloss), y, yhat) = (u = y*yhat; u ≥ 1 ? zero(y) : -y)

struct DWDLoss{T<:Number}
    q::T
end
function (o::DWDLoss)(y, yhat)
    agreement = y * yhat
    q = o.q
    if agreement ≤ q / (q + 1)
        1 - agreement
    else
        (q ^ q / (q + 1) ^ (q + 1)) / agreement ^ q
    end
end
function deriv(loss::DWDLoss, y, yhat)
    u = y * yhat
    q = loss.q
    if u ≤ q / (q + 1)
        -y
    else
       - y * ( q / (q + 1)) ^ (q + 1) / u ^ (q + 1)
    end
end

#-----------------------------------------------------------------------------# Penalties
prox(::typeof(zero), x, s) = x
prox(::typeof(abs2), x, s) = x / (1 + s)
prox(::typeof(abs), x::T, s::Number) where {T} = sign(x) * max(T(0), abs(x) - s)

struct ElasticNet{T}
    α::T
end
(p::ElasticNet)(x) = smooth(abs2(x), abs(x), p.α)
function prox(p::ElasticNet, x, s)
    αs = p.α * s
    prox(abs, x, αs) / (1 + s - αs)
end

#-----------------------------------------------------------------------------# StatLearn
"""
    StatLearn(args...; penalty=zero, rate=LearningRate())

Fit a model (via stochastic approximation) that is linear in the parameters.  The (offline)
objective function that StatLearn approximately minimizes is

``(1/n) ∑ᵢ f(yᵢ, xᵢ'β) + ∑ⱼ λⱼ g(βⱼ),``

where ``fᵢ`` are loss functions of a response variable and linear predictor, ``λⱼ``s are
nonnegative regularization parameters, and ``g`` is a penalty function.

Use `StatLearn` with caution, as stochastic approximation algorithms are inherently noisy.

# Arguments

- `loss = OnlineStats.l2regloss`: The loss function to be (approximately) minimized.
    - Regression Losses:
        - `l2regloss`: Squared error loss
        - `l1regloss`: Absolute error loss
    - Classification (y ∈ {-1, 1}) Losses:
        - `logisticloss`: Logistic regression
        - `l1hingeloss`: Loss function used in Support Vector Machines.
        - `DWDLoss(q)`: Generalized Distance Weighted Discrimination (smoothed `l1hingeloss`)
- `algorithm = SGD()`: The stochastic approximation method to be used.
    - Algorithms based on Stochastic gradient:
        - `SGD()`: Stochastic Gradient Descent
        - `ADAGRAD()`: AdaGrad (adaptive version of SGD)
        - `RMSPROP()`: RMSProp (adaptive version of SGD)
    - Algorithms based on Majorization-Minimization Principle:
        - `MSPI()`: Majorized Stochastic Proximal Iteration
        - `OMAS()`: Online MM via Averaged Surrogate
        - `OMAP()`: Online MM via Averaged Parameter
- `λ = 0.0`: The hyperparameter(s) used for the penalty function
    - User can provide elementwise penalty hyperparameters (`Vector{Float64}`) or single hyperparameter (`Float64`).

## Keyword Arguments

- `penalty`: The regularization function used.
    - `zero`: no penalty (default)
    - `abs`: (LASSO) parameters penalized by their absolute value
    - `abs2`: (Ridge) parameters penalized by their squared value
    - `ElasticNet(α)`: α * (abs penalty) + (1-α) * (abs2 penalty)
- `rate = LearningRate(.6)`

# Example

    x = randn(1000, 5)
    y = x * range(-1, stop=1, length=5) + randn(1000)

    o = fit!(StatLearn(MSPI()), zip(eachrow(x), y))
    coef(o)

    o = fit!(StatLearn(OnlineStats.l1regloss, ADAGRAD()), zip(eachrow(x), y))
    coef(o)
"""
mutable struct StatLearn{A<:Algorithm, L, P, W} <: OnlineStat{XY}
    β::Vector{Float64}
    λ::Vector{Float64}
    gx::Vector{Float64}
    loss::L
    penalty::P
    alg::A
    rate::W
    n::Int
end
function StatLearn(args...; penalty=zero, rate=LearningRate())
    p = 0
    λ = zeros(1)
    loss = l2regloss
    alg = SGD()
    for a in args
        if a isa AbstractVector
            λ = a
        elseif a isa Float64
            λ = fill(a, 1)
        elseif a isa Algorithm
            alg = a
        elseif a isa Integer
            p = a
        elseif hasmethod(a, (Number, Number)) && hasmethod(deriv, (typeof(a), Number, Number))
            loss = a
        else
            @warn """
            Arguments of type $(typeof(a)) are not recognized by StatLearn.  You may be seeing this as
            a result of OnlineStats removing LearnBase, LossFunctions, and PenaltyFunctions as
            dependencies.  See the `StatLearn` docstring for details.

            - Previous behavior for PenaltyFunctions can be restored e.g.

            ```
            OnlineStats.prox(p::Penalty, x, s) = PenaltyFunctions.prox(p, x, s)
            StatLearn(penalty = L1Penalty())
            ```

            - Previous behavior for LossFunctions can be restored e.g.:

            ```
            f(y, yhat) = LossFunctions.value(L2DistLoss(), y, yhat)
            OnlineStats.deriv(::typeof(f), y, yhat) = LossFunctions.deriv(L2DistLoss(), y, yhat)
            StatLearn(f)
            ```
            """
        end
    end
    init!(alg, p)
    StatLearn(zeros(p), λ, zeros(p), loss, penalty, alg, rate, 0)
end

function Base.show(io::IO, o::StatLearn)
    print(io, "StatLearn: ")
    print(io, name(o.alg, false, false))
    o.penalty != zero && print(io, " | mean(λ)=", mean(o.λ))
    print(io, " | ", o.loss)
    print(io, " | ", o.penalty)
    print(io, " | nobs=", nobs(o))
    print(io, " | nvars=", length(o.β))
end
coef(o::StatLearn) = value(o)

function gradient!(o::StatLearn, x, y)
    d_dη = deriv(o.loss, y, predict(o, x))
    for j in eachindex(o.gx)
        o.gx[j] = x[j] * d_dη
    end
end
function _fit!(o::StatLearn{<:SGAlgorithm}, xy)
    x, y = xy
    (o.n += 1) == 1  && init!(o, length(x))
    gradient!(o, x, y)
    update!(o.alg, o.gx)
    updateβ!(o, o.rate(o.n))
end
function init!(o::StatLearn, p)
    o.β = zeros(p)
    o.gx = zeros(p)
    init!(o.alg, p)
    if length(o.λ) == 1
        o.λ = fill(o.λ[1], p)
    end
end

function _merge!(o::StatLearn, o2::StatLearn)
    o.n += o2.n
    γ = nobs(o2) / nobs(o)
    smooth!(o.β, o2.β, γ)
    merge!(o.alg, o2.alg, γ)
    smooth!(o.λ, o2.λ, γ)
end

predict(o::StatLearn, x::VectorOb) = dot(x, o.β)
predict(o::StatLearn, x::AbstractMatrix) = x * o.β
classify(o::StatLearn, x) = sign.(predict(o, x))

function objective(o::StatLearn, x::AbstractMatrix, y::VectorOb)
    mean(o.loss.(y, predict(o,x))) + sum(o.λ .* o.penalty.(o.β))
end

#-----------------------------------------------------------------------# updateβ!
function updateβ!(o::StatLearn{SGD}, γ)
    for j in eachindex(o.β)
        o.β[j] = prox(o.penalty, o.β[j] - γ * o.gx[j], γ * o.λ[j])
    end
end
function updateβ!(o::StatLearn{T}, γ) where {T<:Union{ADAGRAD, RMSPROP}}
    for j in eachindex(o.β)
        s = γ / sqrt(o.alg.h[j] + ϵ)
        o.β[j] = prox(o.penalty, o.β[j] - s * o.gx[j], s * o.λ[j])
    end
end
function updateβ!(o::StatLearn{ADAM}, γ)
    for j in eachindex(o.β)
        s = γ / sqrt(o.alg.v[j] + ϵ)
        o.β[j] = prox(o.penalty, o.β[j] - s * o.alg.m[j], s * o.λ[j])
    end
end
function updateβ!(o::StatLearn{ADAMAX}, γ)
    for j in eachindex(o.β)
        s = γ / ((1 - o.alg.β1^nobs(o)) * o.alg.v[j])
        o.β[j] = prox(o.penalty, o.β[j] - s * o.alg.m[j], s * o.λ[j])
    end
end
function updateβ!(o::StatLearn{ADADELTA}, γ)
    for j in eachindex(o.β)
        s = o.alg.δ[j]
        o.β[j] = prox(o.penalty, o.β[j] - s * o.gx[j], s * o.λ[j])
    end
end

#------------------------------------------------------------------# Majorization-based
# lipschitz_constant (L): f(θ) ≤ f(x) + ∇f(x)'(θ - x) + (L / 2) ||θ - x||^2
lconst(loss, x, y) = error("Unkown Lipschitz constant for loss $loss")

lconst(o::typeof(l2regloss), x, y) = dot(x, x)

lconst(o::typeof(logisticloss), x, y) = .25 * dot(x, x)

lconst(o::DWDLoss, x, y) = (o.q + 1)^2 / o.q * dot(x, x)

#-----------------------------------------------------------------------# OMAS
# L stored in o.alg.a[1]
function _fit!(o::StatLearn{OMAS}, xy)
    x, y = xy
    (o.n += 1) == 1  && init!(o, length(x))
    γ = o.rate(o.n)
    b = o.alg.b
    gradient!(o, x, y)
    ht = lconst(o.loss, x, y)
    L = (o.alg.a[1] = smooth(o.alg.a[1], ht, γ))
    for j in eachindex(o.β)
        b[j] = smooth(b[j], ht * o.β[j] - o.gx[j], γ)
        o.β[j] = prox(o.penalty, b[j] / L, o.λ[j] / L)
    end
end
#-----------------------------------------------------------------------# OMAP
function _fit!(o::StatLearn{OMAP}, xy)
    x, y = xy
    (o.n += 1) == 1 && init!(o, length(x))
    γ = o.rate(o.n)
    gradient!(o, x, y)
    h_inv = inv(lconst(o.loss, x, y))
    for j in eachindex(o.β)
        o.β[j] -= γ * h_inv * o.gx[j]
    end
end
#-----------------------------------------------------------------------# MSPI
function _fit!(o::StatLearn{MSPI}, xy)
    x, y = xy
    (o.n += 1) == 1  && init!(o, length(x))
    γ = o.rate(o.n)
    gradient!(o, x, y)
    γ2 = γ / (1 + γ * lconst(o.loss, x, y))
    for j in eachindex(o.β)
        @inbounds o.β[j] = prox(o.penalty, o.β[j] - γ2 * o.gx[j], γ2 * o.λ[j])
    end
end


text:
# Small sanity tests to ensure changing the rounding of float functions work
using Base.Test

## Float64 checks
# a + b returns a number exactly between prevfloat(1.) and 1., so its
# final result depends strongly on the utilized rounding direction.
a = prevfloat(0.5)
b = 0.5
c = 0x1p-54
d = prevfloat(1.)

# Default rounding direction, RoundNearest
@test a + b === 1.
@test - a - b === -1.
@test a - b === -c
@test b - a === c

# RoundToZero
with_rounding(Float64,RoundToZero) do
    @test a + b === d
    @test - a - b === -d
    @test a - b === -c
    @test b - a === c
end

# Sanity check to see if we have returned to RoundNearest
@test a + b === 1.
@test - a - b === -1.
@test a - b == -c
@test b - a == c

# RoundUp
with_rounding(Float64,RoundUp) do
    @test a + b === 1.
    @test - a - b === -d
    @test a - b === -c
    @test b - a === c
end

# RoundDown
with_rounding(Float64,RoundDown) do
    @test a + b === d
    @test - a - b === -1.
    @test a - b === -c
    @test b - a === c
end

## Float32 checks

a32 = prevfloat(0.5f0)
b32 = 0.5f0
c32 = (1.f0 - prevfloat(1.f0))/2
d32 = prevfloat(1.0f0)

# Default rounding direction, RoundNearest
@test a32 + b32 === 1.0f0
@test - a32 - b32 === -1.0f0
@test a32 - b32 === -c32
@test b32 - a32 === c32

# RoundToZero
with_rounding(Float32,RoundToZero) do
    @test a32 + b32 === d32
    @test - a32 - b32 === -d32
    @test a32 - b32 === -c32
    @test b32 - a32 === c32
end

# Sanity check to see if we have returned to RoundNearest
@test a32 + b32 === 1.0f0
@test - a32 - b32 === -1.0f0
@test a32 - b32 == -c32
@test b32 - a32 == c32

# RoundUp
with_rounding(Float32,RoundUp) do
    @test a32 + b32 === 1.0f0
    @test - a32 - b32 === -d32
    @test a32 - b32 === -c32
    @test b32 - a32 === c32
end

# RoundDown
with_rounding(Float32,RoundDown) do
    @test a32 + b32 === d32
    @test - a32 - b32 === -1.0f0
    @test a32 - b32 === -c32
    @test b32 - a32 === c32
end

max_stars_repo_name: jakebolewski/julia

max_issues_repo_name: jakebolewski/julia

max_forks_repo_name: jakebolewski/julia


text:
include("lqr.jl")
include("two_dimensional_nonlinear_oscillator.jl")


text:
@testset "Test linear solvers" begin
    As = @SMatrix [4..6 -1..1 -1..1 -1..1;-1..1 -6.. -4 -1..1 -1..1;-1..1 -1..1 9..11 -1..1;-1..1 -1..1 -1..1 -11.. -9]
    bs = @SVector [-2..4, 1..8, -4..10, 2..12]

    Am = Matrix(As)
    bm = Vector(bs)
    jac = Jacobi()
    gs = GaussSeidel()
    hbr = HansenBliekRohn()
    kra = LinearKrawczyk()

    for (A, b) in zip([As, Am], [bs, bm])
        xgs = solve(A, b, gs)
        xjac = solve(A, b, jac)
        xhbr = solve(A, b, hbr)
        xkra = solve(A, b, kra)

        @test all(interval_isapprox.(xgs, [-2.6..3.1, -3.9..1.65, -1.48..2.15, -2.35..0.79]; atol=0.01))
        @test all(interval_isapprox.(xjac, [-2.6..3.1, -3.9..1.65, -1.48..2.15, -2.35..0.79]; atol=0.01))

        @test all(interval_isapprox.(xhbr, [-2.5..3.1, -3.9..1.2, -1.4..2.15, -2.35..0.6]; atol=0.01))

        @test all(interval_isapprox.(xkra, [-8..8, -8..8, -8..8, -8..8]; atol=0.01))
    end

    ge = GaussianElimination()
    xge = solve(Am, bm, ge)
    @test all(interval_isapprox.(xge, [-2.6..3.1, -3.9..1.5, -1.43..2.15, -2.35..0.6]; atol=0.01))

    xdef = solve(Am, bm)
    @test all(interval_isapprox.(xdef, [-2.6..3.1, -3.9..1.5, -1.43..2.15, -2.35..0.6]; atol=0.01))

    A = [2..4 -2..1; -1..2 2..4]
    b = [-2..2, -2..2]

    x1 = solve(A, b)
    @test all(interval_isapprox.(x1, [-14..14, -14..14]))
    x2 = solve(A, b, HansenBliekRohn())
    @test all(interval_isapprox.(x2, [-14..14, -14..14]))

    # test exceptions
    @test_throws DimensionMismatch solve(Am, bm[1:end-1])
    @test_throws DimensionMismatch solve(Am[:, 1:end-1], bm, hbr)
    @test_throws DimensionMismatch solve(Am, bm, gs, NoPrecondition(), [1..2, 3..4])
end

@testset "Reduced Row Echelon Form" begin
    A1 = [1..2 1..2;2..2 3..3]
    @test rref(A1) == [2..2 3..3; 0..0 -2..0.5]

    A2 = fill(0..0, 2, 2)
    @test_throws ArgumentError rref(A2)
end


text:
#NetwAnalysis: Network parameter conversion tools
#-------------------------------------------------------------------------------

#=Comments

1) All passive 2-port networks are guaranteed to have S-parameters.  Most others
   cannot represent all passive networks.

   So... indirect conversions should try to use S-parameters as the intermediate.

2) "Default" network parameter conversions are trapped using ::NPT_SConv/NP_SConv
   instead of ::NPType/NetworkParameters.  This allows un-expected conversions
	to be detected, and should avoid infinite recursion issues (ex: continuously
   trying to convert S-parameters -> S-parameters).

3) Network(RT::NPType, np::Network; z0 = nothing) is only a thin interface.
   --> Core algorithm implemented as "_Network" to avoid ambiguity/cyclical issues.
=#

#==Constants
===============================================================================#
const default_z0 = 50.0

#==Convenience types
===============================================================================#

#Network parameter convenience unions:
const NP_SConv = Union{ #Those with S-parameter intermediaries
	ZParameters, YParameters,
	ABCDParameters,
	TParameters
}
const NP_ZConv = Union{ #Those with Z-parameter intermediaries
	HParameters, GParameters,
}

#Network parameter >>type<< convenience unions:
#TODO: auto-calculate from NP_SConv??
const NPT_SConv = Union{ #Those with S-parameter intermediaries
	NPType{:Z}, NPType{:Y},
	NPType{:ABCD},
	NPType{:T}
}
const NPT_ZConv = Union{ #Those with Z-parameter intermediaries
	NPType{:H}, NPType{:G},
}

#==Helper functions
===============================================================================#

#==Data type re-structuring
===============================================================================#

#vector_push: Convert Vector{NwPar{Number}} => NwPar{Vector{Number}}
#-------------------------------------------------------------------------------

#Main algorithm for vector_push
#RT: Result type
function _vector_push(RT::DataType, v::Vector{NT}, z0=nothing) where {NT<:NetworkParameters}
	T = eltype(RT)
	nports = portcount(v[1])
	vlen = length(v)
	result_m = Array{T}(undef, nports, nports)

	for i in 1:vlen
		if portcount(v[i]) != nports
			throw(ArgumentError("Inconsistent number of ports."))
		end
	end

	for row = 1:nports, col = 1:nports
		result_m[row, col] = T(undef, vlen)

		for i in 1:vlen
			result_m[row, col][i] = v[i][row, col]
		end
	end

	if z0 != nothing
		return RT(z0, result_m)
	else
		return RT(result_m)
	end
end
vector_push(np::Vector{NetworkParametersRef{TP, NP, T}}) where {TP, NP, T} =
	_vector_push(NetworkParametersRef{TP, NP, Vector{T}}, np, z0=np.z0)
vector_push(np::Vector{NetworkParametersNoRef{TP, NP, T}}) where {TP, NP, T} =
	_vector_push(NetworkParametersNoRef{TP, NP, Vector{T}}, np)

#vector_pull: Convert NwPar{Vector{Number}} => Vector{NwPar{Number}}
#-------------------------------------------------------------------------------

#Main algorithm for vector_pull:
#ET: Element type for result
function _vector_pull(ET::DataType, np::NetworkParameters, z0=nothing)
	nports = portcount(np)
	vlen = length(np[1,1])
	result = Array{ET}(undef, vlen)

	for row = 1:nports, col = 1:nports
		if length(np[row, col]) != vlen
			throw(ArgumentError("Inconsistent vector lengths."))
		end
	end

	newelem_z0() = ET(z0, NetworkParameterMatrix{eltype(ET)}(undef, nports, nports))
	newelem_noref() = ET(NetworkParameterMatrix{eltype(ET)}(undef, nports, nports))
	newelem = z0 != nothing ? newelem_z0 : newelem_noref

	for i in 1:vlen
		inner = newelem() #Inner matrix

		for row = 1:nports, col = 1:nports
			inner.m[row, col] = np[row, col][i]
		end
		result[i] = inner
	end
	return result
end

vector_pull(np::NetworkParametersRef{TP, NP, Vector{T}}) where {TP, NP, T} =
	_vector_pull(NetworkParametersRef{TP, NP, T}, np, z0=np.z0)
vector_pull(np::NetworkParametersNoRef{TP, NP, Vector{T}}) where {TP, NP, T} =
	_vector_pull(NetworkParametersNoRef{TP, NP, T}, np)

#submatrix: Obtain subset of a NetworkParameters matrix
#-------------------------------------------------------------------------------
function submatrix(np::NetworkParameters, ports::Vector{PT}=Int[]) where {PT<:Integer}
	nports = length(ports)

	if !in(:Y, Set([:S, :Z, :Y]))
		throw(ArgumentError("Can only compute subset of S, Z, or Y matrices"))
	end

	ET = eltype(np)
	result_m = NetworkParameterMatrix{ET}(undef, nports, nports)

	for row=1:nports, col=1:nports
		result_m[row, col] = np[ports[row], ports[col]]
	end

	if typeof(np)<:NetworkParametersRef
		return Network(NPType(np), result_m, z0=np.z0)
	else
		return Network(NPType(np), result_m)
	end

end

#==Reference impedance transformation
===============================================================================#

#Generate T-matrix of an impedance transformer:
function z0xfrm_matrix(z01::Float64, z02::Float64)
	#Forward reflection coefficient:
	ΓF = (z02 - z01) / (z02+z01)

	return ((1 / (1+ΓF)) * sqrt(z02/z01)) .* [1 ΓF; ΓF 1]
end

#T-parameter impedance transformation:
function z0xfrm(T::TParameters, z0::Float64)
	#NOTE:
	#Impedance transformation done by sandwiching current T-matrix between 2
	#impedance transformers: [z01|z02] [T] [z02|z01]
	#where z01 is the ref impedance of the returned T-matrix.

	T1 = z0xfrm_matrix(z0, T.z0)
	T2 = z0xfrm_matrix(z0, T.z0)
	result = T.m

	return Network(NPType(:T), T1 * T.m * T2, z0=z0)
end

#==Network parameter conversions: {T,ABCD,Z,Y}=>S
===============================================================================#

#Convert T => S-parameters (equations always maintains z0):
function npconv(RT::NPType{:S}, T::TParameters)
	(t11, t12, t21, t22) = mx2elem(T)
	s11 = t21/t11
	s12 = t22 - t21*t12/t11
	s21 = 1/t11
	s22 = -t12/t11
	return Network(RT, [s11 s12; s21 s22], z0=T.z0)
end

#Convert ABCD => S-parameters:
function npconv(RT::NPType{:S}, ABCD::ABCDParameters, z0::Float64)
	(A, B, C, D) = mx2elem(ABCD)
	B_z0 = B / z0; Cxz0 = C * z0
	denom = A + B_z0 + Cxz0 + D
	s11 = (A + B_z0 - Cxz0 - D) / denom
	s12 = (2(A*D - B*C)) / denom
	s21 = 2 / denom
	s22 = (-A + B_z0 - Cxz0 + D) / denom
	return Network(RT, [s11 s12; s21 s22], z0=z0)
end

#Convert Z => S-parameters:
function npconv(RT::NPType{:S}, Z::ZParameters, z0::Float64)
	(z11, z12, z21, z22) = mx2elem(Z)
	z12z21 = z12*z21
	denom = (z11+z0)*(z22+z0)-z12z21
	s11 = ((z11-z0)*(z22+z0) - z12z21) / denom
	s12 = 2z12*z0 / denom
	s21 = 2z21*z0 / denom
	s22 = ((z11+z0)*(z22-z0) - z12z21) / denom
	return Network(RT, [s11 s12; s21 s22], z0=z0)
end

#Convert Y => S-parameters:
function npconv(RT::NPType{:S}, Y::YParameters, z0::Float64)
	(y11, y12, y21, y22) = mx2elem(Y)
	y0 = 1/z0
	y12y21 = y12*y21
	denom = (y11+y0)*(y22+y0) - y12y21
	s11 = ((y0-y11)*(y0+y22) + y12y21) / denom
	s12 = -2y12*y0 / denom
	s21 = -2y21*y0 / denom
	s22 = ((y0+y11)*(y0-y22) + y12y21) / denom
	return Network(RT, [s11 s12; s21 s22], z0=z0)
end

#==Network parameter conversions: S=>{T,ABCD,Z,Y}
===============================================================================#

#Convert 2-port S-parameters => T (equations always maintains z0):
function npconv(RT::NPType{:T}, S::SParameters{2})
	(s11, s12, s21, s22) = mx2elem(S)
	t11 = 1/s21
	t12 = -s22/s21
	t21 = s11/s21
	t22 = s12 - s11*s22/s21
	return Network(RT, [t11 t12; t21 t22], z0=S.z0)
end

#Convert 2-port S-parameters => ABCD:
function npconv(RT::NPType{:ABCD}, S::SParameters{2})
	(s11, s12, s21, s22) = mx2elem(S)
	z0 = S.z0
	s12_2 = s12 / 2
	s21x2 = 2*s21
	A = ((1+s11)*(1-s22)) / s21x2 + s12_2
	B = z0*( ((1+s11)*(1+s22)) / s21x2 - s12_2)
	C = (((1-s11)*(1-s22)) / s21x2 - s12_2) / z0
	D = ((1-s11)*(1+s22)) / s21x2 + s12_2
	return Network(RT, [A B; C D])
end

#Convert 2-port S-parameters => Z:
#TODO: Arbitrary number of ports
function npconv(RT::NPType{:Z}, S::SParameters{2})
	(s11, s12, s21, s22) = mx2elem(S)
	z0 = S.z0
	s12s21 = s12*s21
	cFact = z0 / ((1-s11)*(1-s22) - s12s21)
	z11 = cFact * ((1+s11)*(1-s22) + s12s21)
	z12 = cFact * 2 * s12
	z21 = cFact * 2 * s21
	z22 = cFact * ((1-s11)*(1+s22) + s12s21)
	return Network(RT, [z11 z12; z21 z22])
end

#Convert 2-port S-parameters => Y:
#TODO: Arbitrary number of ports
function npconv(RT::NPType{:Y}, S::SParameters{2})
	(s11, s12, s21, s22) = mx2elem(S)
	z0 = S.z0
	s12s21 = s12*s21
	cFact = 1 / (z0 * ((1+s11)*(1+s22) - s12s21))
	y11 = cFact * ((1-s11)*(1+s22) + s12s21)
	y12 = -2 * s12 * cFact
	y21 = -2 * s21 * cFact
	y22 = cFact * ((1+s11)*(1-s22) + s12s21)
	return Network(RT, [y11 y12; y21 y22])
end

#==Network parameter conversions: {G, H} <=> Z
===============================================================================#

#Convert H-parameters => Z:
function npconv(RT::NPType{:Z}, H::HParameters)
	(h11, h12, h21, h22) = mx2elem(H)
	z11 = (h11*h22 - h12*h21) / h22
	z12 = h12 / h22
	z21 = -h21 / h22
	z22 = 1 / h22
	return Network(RT, [z11 z12; z21 z22])
end

#Convert G-parameters => Z:
function npconv(RT::NPType{:Z}, G::GParameters)
	(g11, g12, g21, g22) = mx2elem(G)
	z11 = 1 / g11
	z12 = -g12 / g11
	z21 = g21 / g11
	z22 = (g11*g22 - g12*g21) / g11
	return Network(RT, [z11 z12; z21 z22])
end

#Convert 2-port Z-parameters => H:
function npconv(RT::NPType{:H}, Z::ZParameters{2})
	(z11, z12, z21, z22) = mx2elem(Z)
	h11 = (z11*z22 - z12*z21) / z22
	h12 = z12 / z22
	h21 = -z21 / z22
	h22 = 1 / z22
	return Network(RT, [h11 h12; h21 h22])
end

#Convert 2-port Z-parameters => G:
function npconv(RT::NPType{:G}, Z::ZParameters{2})
	(z11, z12, z21, z22) = mx2elem(Z)
	g11 = 1 / z11
	g12 = -z12 / z11
	g21 = z21 / z11
	g22 = (z11*z22 - z12*z21) / z11
	return Network(RT, [g11 g12; g21 g22])
end

#==Reference impedance transformation/passthrough
===============================================================================#
z0xfrm(np::Network, z0::Float64) where TP =
	error("Unable to perform z0 transformation for $(typeof(np)).")
#S-parameter impedance transformation:
z0xfrm(np::SParameters, z0::Float64) =
	npconv(NPType(:S), z0xfrm(npconv(NPType(:T), np), z0)) #TODO: implement directly

#==Network parameter conversions: S <=> T
===============================================================================#
#S -> T WITH z0 conversion - TODO: implement directly
npconv(RT::NPType{:T}, np::SParameters, z0::Float64) =
	z0xfrm(npconv(RT, np), z0)
#T -> S WITH z0 conversion - TODO: implement directly
npconv(RT::NPType{:S}, np::TParameters, z0::Float64) =
	npconv(RT, z0xfrm(np, z0))

#==Basic traps for npconv()
===============================================================================#
#Catch-all:
npconv(npt::NPType, np::Network) =
	error("Unable to convert to $(typeof(np)) -> :$(Symbol(npt)).")
npconv(npt::NPType, np::Network, z0::Float64) =
	error("Unable to convert to $(typeof(np)) -> :$(Symbol(npt)), z0=$z0.")
#TODO: copy(np) on conversion????
npconv(::NPType{TP}, np::NetworkParameters{TP}) where {TP} = np #Already correct format
#Traps to z0 transformation/leaving parameters untouched:
npconv(::NPType{TP}, np::NetworkParameters{TP}, z0::Float64) where {TP} =
	((z0 != np.z0) ? z0xfrm(np, z0) : np) #Might require impedance transformation.

#==Indirect network parameter conversions
===============================================================================#
_npconv(args...) = npconv(args...) #Default: Use direct conversion

for dest in QuoteNode.([:ABCD, :Y, :Z]); for src in QuoteNode.([:ABCD, :Y, :Z]);
if src == dest; continue; end; @eval begin #CODEGEN
	_npconv(RT::NPType{$dest}, np::NetworkParameters{$src}) =
		npconv(RT, npconv(NPType(:S), np, default_z0))
end; end; end #CODEGEN----------------------------------------------------------
for dest in QuoteNode.([:ABCD, :Y, :Z]); for src in QuoteNode.([:T]); @eval begin #CODEGEN
	_npconv(RT::NPType{$dest}, np::NetworkParameters{$src}) =
		npconv(RT, npconv(NPType(:S), np)) #No z0 in T-> S xform
end; end; end #CODEGEN----------------------------------------------------------
for dest in QuoteNode.([:T]); for src in QuoteNode.([:ABCD, :Y, :Z]); @eval begin #CODEGEN
	_npconv(RT::NPType{$dest}, np::NetworkParameters{$src}, z0) =
		npconv(RT, npconv(NPType(:S), np, z0), z0)
end; end; end #CODEGEN----------------------------------------------------------

for dest in QuoteNode.([:G, :H]); for src in QuoteNode.([:ABCD, :Y, :Z]); @eval begin #CODEGEN
	_npconv(RT::NPType{$dest}, np::NetworkParameters{$src}) =
		npconv(RT, npconv(NPType(:Z), npconv(NPType(:S), np, default_z0)))
end; end; end #CODEGEN----------------------------------------------------------
for dest in QuoteNode.([:G, :H]); for src in QuoteNode.([:T, :S]); @eval begin #CODEGEN
	_npconv(RT::NPType{$dest}, np::NetworkParameters{$src}) =
		npconv(RT, npconv(NPType(:Z), npconv(NPType(:S), np))) #No z0 in T-> S xform
end; end; end #CODEGEN----------------------------------------------------------
for dest in QuoteNode.([:ABCD, :Y, :Z]); for src in QuoteNode.([:G, :H]); @eval begin #CODEGEN
	_npconv(RT::NPType{$dest}, np::NetworkParameters{$src}) =
		npconv(RT, npconv(NPType(:S), npconv(NPType(:Z), np), default_z0))
end; end; end #CODEGEN----------------------------------------------------------
for dest in QuoteNode.([:T, :S]); for src in QuoteNode.([:G, :H]); @eval begin #CODEGEN
	_npconv(RT::NPType{$dest}, np::NetworkParameters{$src}, z0) =
		npconv(RT, npconv(NPType(:S), npconv(NPType(:Z), np), z0))
end; end; end #CODEGEN----------------------------------------------------------

#==High-level interface (z0 optional parameter):
===============================================================================#
Network(RT::NPType, np::Network) = _npconv(RT, np)
#Default: maintain z0 between scattering parameters:
Network(RT::NPType{:S}, np::NetworkParametersRef) = _npconv(RT, np, np.z0) #Maintain z0.
Network(RT::NPType{:T}, np::NetworkParametersRef) = _npconv(RT, np, np.z0) #Maintain z0.
#Allow user overwrite/use default_z0 when converting to scattering parameters:
Network(RT::NPType{:S}, np::Network; z0::Real = default_z0) = _npconv(RT, np, Float64(z0))
Network(RT::NPType{:T}, np::Network; z0::Real = default_z0) = _npconv(RT, np, Float64(z0))
#TODO: Deal with complex z0??

#Last line


text:
using FastGaussQuadrature
import GSL

"""
    U = upsampling_matrix(n, m)

um = U*un upsamples a function u(x) from n to m GL nodes on [-1,1] \\
Can be used to create a downsampling matrix as well
"""
function upsampling_matrix(n, m)
    L = legendre_matrix(n)
    B = Array{Float64}(m, n)
    glpoints, glweights = gausslegendre(m)
    for i=1:m
        P = GSL.sf_legendre_Pl_array(n-1, glpoints[i])
        for j=1:n
            B[i,j] = P[j]
        end
    end
    return B*L
end

"""
    L = legendre_matrix(order)

L computes Legendre expansion coefficients for a function
defined at the Gauss-Legendre quadrature nodes on [-1, 1] \\
``c = L*f  ->  c_l = (2l+1)/2 \\sum_n P_l(x_n) f_n w_n``
"""
function legendre_matrix(n)
    glpoints, glweights = gausslegendre(n)
    L = Array{Float64}(n, n)
    for i=1:n
        P = GSL.sf_legendre_Pl_array(n-1, glpoints[i])
        for j=1:n
            l = j-1
            L[j, i] = P[j]*glweights[i]*(2*l+1)/2
        end
    end
    return L
end

"""
    D = legendre_diff_matrix(order)

D approximates the derivative of a function defined at
the Gauss-Legendre quadrature nodes on [-1, 1] \\
``D*f =~ f'``
"""
function legendre_diff_matrix(n)
    glpoints, glweights = gausslegendre(n)
    dPt = Array{Float64}(n, n)
    for i=1:n
        _, dP = legendre_Pl_deriv_array(n-1, glpoints[i])
        for j=1:n
            dPt[i, j] = dP[j]
        end
    end
    L = legendre_matrix(n)
    return dPt*L
end

"""
P, D = legendre_Pl_deriv_array(n, x) \\
Recurrence relation used: [http://dlmf.nist.gov/18.9]

As fast as the GSL implementation, but also takes x outside [-1, 1].
    """
function legendre_Pl_deriv_array(n, x)
    P = Array{Complex{Float64}}(n+1)
    D = Array{Complex{Float64}}(n+1)
    P[1] = 1.0 # l=0
    D[1] = 0.0
    P[2] = x # l=1
    D[2] = 1.0
    for l=1:n-1
        # Compute l+1
        P[l+1+1] = ( (2*l+1)*x * P[l+1] - l*P[l-1+1] ) / (l+1)
        D[l+1+1] = ( (2*l+1)*(P[l+1] + x * D[l+1]) - l*D[l-1+1] ) / (l+1)
    end
    P, D
end

"""
    t, relres, iter = newton_legendre(c, zr, t0, maxiter, tol)

Find root t of Legendre expansion c such that
``\\sum_n c_n*P_n(t) = zr``
"""
function newton_legendre(c, zr, t0, maxiter, tol)
    t = t0
    iter = 0
    dt = Inf
    lmax = length(c)-1
    relres = tol
    for iter=1:maxiter
        P, D = legendre_Pl_deriv_array(lmax, t)
        f = sum(P.*c) - zr
        fp = sum(D.*c)
        dt = -f / fp
        t = t + dt
        relres = abs(dt/t)
        if relres < tol
            break
        end
    end
    return t, relres, iter
end


text:
## old code from Maxwell.jl

########### PWE (k-space) grid initialization routines #################

function k_mn(k::SArray{Tuple{3},Float64,1,3})
    mag = sqrt(k[1]^2 + k[2]^2 + k[3]^2)
    if mag==0
        n = SVector(0.,1.,0.)
        m = SVector(0.,0.,1.)
    else
        if k[1]==0. && k[2]==0.    # put n in the y direction if k+G is in z
            n = SVector(0.,1.,0.)
        else                                # otherwise, let n = z x (k+G), normalized
            temp = SVector(0.,0.,1.) × k
            n = temp / sqrt( temp[1]^2 + temp[2]^2 + temp[3]^2 )
        end
    end

    # m = n x (k+G), normalized
    mtemp = n × k
    m = mtemp / sqrt( mtemp[1]^2 + mtemp[2]^2 + mtemp[3]^2 )
    return m,n
end

struct KVec
    k::SVector{3,Float64}   # 3-vector
    mag::Float64            # vector magnitude
    m::SVector{3,Float64}   # k vector normal with even parity in y
    n::SVector{3,Float64}   # k vector normal with even parity in x
end

KVec(k::SVector{3,Float64}) = KVec(
    k,
    sqrt(sum(abs2(ki) for ki in k)), # sqrt(sum(ki^2 for ki in k)),
    k_mn(k)...,
)

function kpG(k::SVector{3,Float64},g::MaxwellGrid)::Array{KVec,3}
    [KVec(k-gx-gy-gz) for gx=g.gx, gy=g.gy, gz=g.gz]
    # [KVec(k+gx+gy+gz) for gx=g.gx, gy=g.gy, gz=g.gz]
end

function kpG(kz::Float64,g::MaxwellGrid)::Array{KVec,3}
    # kpG(SVector(0.,0.,kz),g::MaxwellGrid)::Array{KVec,3}
    [KVec(SVector(0.,0.,kz)-gx-gy-gz) for gx=g.gx, gy=g.gy, gz=g.gz]
end

########### non-mutating operatators using StaticArrays #################

function KplusG(kz,gx,gy,gz)
	# scale = ds.kpG[i,j,k].mag
	kpg = [-gx; -gy; kz-gz]
	mag = norm(kpg) #sqrt(sum(abs2.(kpg)))
	if mag==0
		n = [0.; 1.; 0.] # SVector(0.,1.,0.)
		m = [0.; 0.; 1.] # SVector(0.,0.,1.)
	else
		if kpg[1]==0. && kpg[2]==0.    # put n in the y direction if k+G is in z
			n = [0.; 1.; 0.] #SVector(0.,1.,0.)
		else                                # otherwise, let n = z x (k+G), normalized
			ntemp = [0.; 0.; 1.] × kpg  #SVector(0.,0.,1.) × kpg
			n = ntemp / norm(ntemp) # sqrt(sum(abs2.(ntemp)))
		end
	end
	# m = n x (k+G), normalized
	mtemp = n × kpg
	m = mtemp / norm(mtemp) # sqrt(sum(abs2.(mtemp))) #sqrt( mtemp[1]^2 + mtemp[2]^2 + mtemp[3]^2 )
	return kpg, mag, m, n
end

function t2c(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
		Hout[1,i,j,k] = ( Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * n[1] ) * mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * m[2] + Hin[2,i,j,k] * n[2] ) * mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * m[3] + Hin[2,i,j,k] * n[3] ) * mag
    end
    return Hout
end

function c2t(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,2,Nx,Ny,Nz) # Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] =  Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        Hout[2,i,j,k] =  Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
    end
    return Hout
end

function zcross_t2c(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] = -Hin[1,i,j,k] * m[2] - Hin[2,i,j,k] * n[2]
        Hout[2,i,j,k] =  Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * n[1]
        Hout[3,i,j,k] = 0.0
    end
    return Hout
end

function kcross_t2c(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] = ( Hin[1,i,j,k] * n[1] - Hin[2,i,j,k] * m[1] ) * -mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * n[2] - Hin[2,i,j,k] * m[2] ) * -mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * n[3] - Hin[2,i,j,k] * m[3] ) * -mag
    end
    return Hout
end

function kcross_c2t(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,2,Nx,Ny,Nz) # Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        at1 = Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        at2 = Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
        Hout[1,i,j,k] =  -at2 * mag
        Hout[2,i,j,k] =  at1 * mag
    end
    return Hout
end

function kcrossinv_t2c(Hin,kz,gx,gy,gz)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        Hout[1,i,j,k] = ( Hin[1,i,j,k] * n[1] - Hin[2,i,j,k] * m[1] ) / mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * n[2] - Hin[2,i,j,k] * m[2] ) / mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * n[3] - Hin[2,i,j,k] * m[3] ) / mag
    end
    return Hout
end

function kcrossinv_c2t(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,2,Nx,Ny,Nz) # Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = KplusG(kz,gx[i],gy[j],gz[k])
        scale = -1 / mag
        at1 = Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        at2 = Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
        Hout[1,i,j,k] =  at2 / mag
        Hout[2,i,j,k] =  -at1 / mag
    end
    return Hout
end

function ε⁻¹_dot(Hin,ε⁻¹)
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        Hout[1,i,j,k] =  ε⁻¹[1,1,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,1,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,1,i,j,k]*Hin[3,i,j,k]
        Hout[2,i,j,k] =  ε⁻¹[1,2,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,2,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,2,i,j,k]*Hin[3,i,j,k]
        Hout[3,i,j,k] =  ε⁻¹[1,3,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,3,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,3,i,j,k]*Hin[3,i,j,k]
        # Hout[1,i,j,k] =  ε⁻¹[i,j,k][1,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][1,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][1,3]*Hin[3,i,j,k]
        # Hout[2,i,j,k] =  ε⁻¹[i,j,k][2,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][2,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][2,3]*Hin[3,i,j,k]
        # Hout[3,i,j,k] =  ε⁻¹[i,j,k][3,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][3,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][3,3]*Hin[3,i,j,k]
    end
    return Hout
end

function ε_dot_approx(Hin,ε⁻¹)
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = similar(Hin,3,Nx,Ny,Nz) # Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        ε_ave = 3 / tr(ε⁻¹[:,:,i,j,k])
        Hout[1,i,j,k] =  ε_ave * Hin[1,i,j,k]
        Hout[2,i,j,k] =  ε_ave * Hin[2,i,j,k]
        Hout[3,i,j,k] =  ε_ave * Hin[3,i,j,k]
    end
    return Hout
end

function M(Hin::AbstractArray{T,4},ε⁻¹,kz,gx,gy,gz)::AbstractArray{T,4} where T
    d = fft(kcross_t2c(Hin,kz,gx,gy,gz),(2:4));
    e = ifft(ε⁻¹_dot(d,ε⁻¹),(2:4)); # (-1/(π)) .*
    kcross_c2t(e,kz,gx,gy,gz)
end

function Mₖ(Hin,ε⁻¹,kz,gx,gy,gz)
    d = fft(zcross_t2c(Hin,kz,gx,gy,gz),(2:4));
    e = ifft(ε⁻¹_dot(d,ε⁻¹),(2:4)); # (-1/(π)) .*
    kcross_c2t(e,kz,gx,gy,gz)
end

function M(Hin::AbstractArray{T,1},ε⁻¹,kz,gx,gy,gz)::AbstractArray{T,1} where T
    Nx,Ny,Nz = size(ε⁻¹)[end-2:end]
    HinA = reshape(Hin,(2,Nx,Ny,Nz))
    HoutA = M(HinA,ε⁻¹,kz,gx,gy,gz)
    return vec(HoutA)
end

# M̂(ε⁻¹,kz,gx,gy,gz) = LinearMap{T}(H -> M(H::AbstractArray{T,1},ε⁻¹,kz,gx,gy,gz),(2*length(gx)*length(gy)*length(gz)),ishermitian=true,ismutating=false)
# M̂(ε⁻¹,kz,gx,gy,gz) = LinearMap{T}(H -> M(H::AbstractArray{T,1},ε⁻¹,kz,gx,gy,gz),(2*length(gx)*length(gy)*length(gz)),ishermitian=true,ismutating=false) where T
M̂(ε⁻¹,kz,gx,gy,gz) = LinearMap(H -> M(H,ε⁻¹,kz,gx,gy,gz),(2*length(gx)*length(gy)*length(gz)),ishermitian=true,ismutating=false)

function Mₖ(Hin::AbstractArray{T,1},ε⁻¹,kz,gx,gy,gz)::AbstractArray{T,1} where T
    Nx,Ny,Nz = size(ε⁻¹)[end-2:end]
    HinA = reshape(Hin,(2,Nx,Ny,Nz))
    HoutA = Mₖ(HinA,ε⁻¹,kz,gx,gy,gz)
    return vec(HoutA)
end

# M̂ₖ(ε⁻¹,kz,gx,gy,gz) = LinearMap{T}(H -> Mₖ(H::AbstractArray{T,1},ε⁻¹,kz,gx,gy,gz),(2*length(gx)*length(gy)*length(gz)),ishermitian=true,ismutating=false) where T
M̂ₖ(ε⁻¹,kz,gx,gy,gz) = LinearMap(H -> Mₖ(H,ε⁻¹,kz,gx,gy,gz),(2*length(gx)*length(gy)*length(gz)),ishermitian=true,ismutating=false)
# M̂ₖ(ε⁻¹,kz,gx,gy,gz) = LinearMap{T}(H -> Mₖ(H::AbstractArray{T,1},ε⁻¹,kz,gx,gy,gz),(2*length(gx)*length(gy)*length(gz)),ishermitian=true,ismutating=false)

#

# Zygote versions

function sum2(op,arr)
    return sum(op,arr)
end

function sum2adj( Δ, op, arr )
    n = length(arr)
    g = x->Δ*Zygote.gradient(op,x)[1]
    return ( nothing, map(g,arr))
end

Zygote.@adjoint function sum2(op,arr)
    return sum2(op,arr),Δ->sum2adj(Δ,op,arr)
end

function zyg_KplusG(kz,gx,gy,gz)
	# scale = ds.kpG[i,j,k].mag
	kpg = [-gx; -gy; kz-gz]
	mag = sqrt(sum2(abs2,kpg))
	if mag==0
		n = [0.; 1.; 0.] # SVector(0.,1.,0.)
		m = [0.; 0.; 1.] # SVector(0.,0.,1.)
	else
		if kpg[1]==0. && kpg[2]==0.    # put n in the y direction if k+G is in z
			n = [0.; 1.; 0.] #SVector(0.,1.,0.)
		else                                # otherwise, let n = z x (k+G), normalized
			ntemp = [0.; 0.; 1.] × kpg  #SVector(0.,0.,1.) × kpg
			n = ntemp / sqrt(sum2(abs2,ntemp))
		end
	end
	# m = n x (k+G), normalized
	mtemp = n × kpg
	m = mtemp / sqrt(sum2(abs2,mtemp)) #sqrt( mtemp[1]^2 + mtemp[2]^2 + mtemp[3]^2 )
	return kpg, mag, m, n
end

function zyg_t2c(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(3,size(Hin)[2:end]...))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = Zygote.forwarddiff(kz) do kz
                            zyg_KplusG(kz,gx[i],gy[j],gz[k])
                        end
		Hout[1,i,j,k] = ( Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * n[1] ) * mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * m[2] + Hin[2,i,j,k] * n[2] ) * mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * m[3] + Hin[2,i,j,k] * n[3] ) * mag
    end
    return copy(Hout)
end

function zyg_c2t(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = Zygote.forwarddiff(kz) do kz
                            zyg_KplusG(kz,gx[i],gy[j],gz[k])
                        end
        Hout[1,i,j,k] =  Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        Hout[2,i,j,k] =  Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
    end
    return copy(Hout)
end

function zyg_zcross_t2c(Hin,kz,gx,gy,gz)
    # Hout = zeros(ComplexF64,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = Zygote.forwarddiff(kz) do kz
                            zyg_KplusG(kz,gx[i],gy[j],gz[k])
                        end
        Hout[1,i,j,k] = -Hin[1,i,j,k] * m[2] - Hin[2,i,j,k] * n[2]
        Hout[2,i,j,k] =  Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * n[1]
        Hout[3,i,j,k] = 0.0
    end
    return copy(Hout)
end

function zyg_kcross_t2c(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = Zygote.forwarddiff(kz) do kz
                            zyg_KplusG(kz,gx[i],gy[j],gz[k])
                        end
        Hout[1,i,j,k] = ( Hin[1,i,j,k] * n[1] - Hin[2,i,j,k] * m[1] ) * -mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * n[2] - Hin[2,i,j,k] * m[2] ) * -mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * n[3] - Hin[2,i,j,k] * m[3] ) * -mag
    end
    return copy(Hout)
end

function zyg_kcross_c2t(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = Zygote.forwarddiff(kz) do kz
                            zyg_KplusG(kz,gx[i],gy[j],gz[k])
                        end
        at1 = Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        at2 = Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
        Hout[1,i,j,k] =  -at2 * mag
        Hout[2,i,j,k] =  at1 * mag
    end
    return copy(Hout)
end

function zyg_kcrossinv_t2c(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = Zygote.forwarddiff(kz) do kz
                            zyg_KplusG(kz,gx[i],gy[j],gz[k])
                        end
        Hout[1,i,j,k] = ( Hin[1,i,j,k] * n[1] - Hin[2,i,j,k] * m[1] ) / mag
        Hout[2,i,j,k] = ( Hin[1,i,j,k] * n[2] - Hin[2,i,j,k] * m[2] ) / mag
        Hout[3,i,j,k] = ( Hin[1,i,j,k] * n[3] - Hin[2,i,j,k] * m[3] ) / mag
    end
    return copy(Hout)
end

function zyg_kcrossinv_c2t(Hin,kz,gx,gy,gz)
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
	@inbounds for i=1:Nx,j=1:Ny,k=1:Nz
		kpg, mag, m, n = Zygote.forwarddiff(kz) do kz
                            zyg_KplusG(kz,gx[i],gy[j],gz[k])
                        end
        scale = -1 / mag
        at1 = Hin[1,i,j,k] * m[1] + Hin[2,i,j,k] * m[2] + Hin[3,i,j,k] * m[3]
        at2 = Hin[1,i,j,k] * n[1] + Hin[2,i,j,k] * n[2] + Hin[3,i,j,k] * n[3]
        Hout[1,i,j,k] =  at2 / mag
        Hout[2,i,j,k] =  -at1 / mag
    end
    return copy(Hout)
end

function zyg_ε⁻¹_dot(Hin,ε⁻¹)
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        Hout[1,i,j,k] =  ε⁻¹[1,1,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,1,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,1,i,j,k]*Hin[3,i,j,k]
        Hout[2,i,j,k] =  ε⁻¹[1,2,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,2,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,2,i,j,k]*Hin[3,i,j,k]
        Hout[3,i,j,k] =  ε⁻¹[1,3,i,j,k]*Hin[1,i,j,k] + ε⁻¹[2,3,i,j,k]*Hin[2,i,j,k] + ε⁻¹[3,3,i,j,k]*Hin[3,i,j,k]
        # Hout[1,i,j,k] =  ε⁻¹[i,j,k][1,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][1,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][1,3]*Hin[3,i,j,k]
        # Hout[2,i,j,k] =  ε⁻¹[i,j,k][2,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][2,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][2,3]*Hin[3,i,j,k]
        # Hout[3,i,j,k] =  ε⁻¹[i,j,k][3,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][3,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][3,3]*Hin[3,i,j,k]
    end
    return copy(Hout)
end

function zyg_ε_dot_approx(Hin,ε⁻¹)
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    @inbounds for i=1:Nx,j=1:Ny,k=1:Nz
        ε_ave = 3 / tr(ε⁻¹[:,:,i,j,k])
        Hout[1,i,j,k] =  ε_ave * Hin[1,i,j,k]
        Hout[2,i,j,k] =  ε_ave * Hin[2,i,j,k]
        Hout[3,i,j,k] =  ε_ave * Hin[3,i,j,k]
    end
    return copy(Hout)
end

function zyg_M(Hin,ε⁻¹,kz,gx,gy,gz)
    d = fft(zyg_kcross_t2c(Hin,kz,gx,gy,gz),(2:4));
    e = ifft(zyg_ε⁻¹_dot(d,ε⁻¹),(2:4)); # (-1/(π)) .*
    zyg_kcross_c2t(e,kz,gx,gy,gz)
end

# function zyg_M(Hin::AbstractArray{ComplexF64,1},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ds::MaxwellData)::Array{ComplexF64,1}
#     HinA = reshape(Hin,(2,ds.Nx,ds.Ny,ds.Nz))
#     HoutA = M(HinA,ε⁻¹,ds)
#     return vec(HoutA)
# end

function zyg_Mₖ(Hin,ε⁻¹,kz,gx,gy,gz)
    d = fft(zyg_zcross_t2c(Hin,kz,gx,gy,gz),(2:4));
    e = ifft(zyg_ε⁻¹_dot(d,ε⁻¹),(2:4));
    zyg_kcross_c2t(e,kz,gx,gy,gz)
end

# old non-mutating operator methods using StaticArrays

function t2c(Hin::AbstractArray{ComplexF64,4},kpG::AbstractArray{KVec,3})::AbstractArray{ComplexF64,4}
    # Hout = Array{ComplexF64}(undef,(3,size(Hin)[2:end]...))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    for i=1:Nx,j=1:Ny,k=1:Nz
        @inbounds scale = ds.kpG[i,j,k].mag
        @inbounds Hout[1,i,j,k] = ( Hin[1,i,j,k] * kpG[i,j,k].m[1] + Hin[2,i,j,k] * kpG[i,j,k].n[1] ) * scale
        @inbounds Hout[2,i,j,k] = ( Hin[1,i,j,k] * kpG[i,j,k].m[2] + Hin[2,i,j,k] * kpG[i,j,k].n[2] ) * scale
        @inbounds Hout[3,i,j,k] = ( Hin[1,i,j,k] * kpG[i,j,k].m[3] + Hin[2,i,j,k] * kpG[i,j,k].n[3] ) * scale
    end
    return copy(Hout)
end

t2c(Hin::AbstractArray{ComplexF64,4},ds::MaxwellData)::AbstractArray{ComplexF64,4} = t2c(Hin,ds.kpG)

function c2t(Hin::AbstractArray{ComplexF64,4},kpG::AbstractArray{KVec,3})::AbstractArray{ComplexF64,4}
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
    for i=1:Nx,j=1:Ny,k=1:Nz
        @inbounds Hout[1,i,j,k] =  Hin[1,i,j,k] * kpG[i,j,k].m[1] + Hin[2,i,j,k] * kpG[i,j,k].m[2] + Hin[3,i,j,k] * kpG[i,j,k].m[3]
        @inbounds Hout[2,i,j,k] =  Hin[1,i,j,k] * kpG[i,j,k].n[1] + Hin[2,i,j,k] * kpG[i,j,k].n[2] + Hin[3,i,j,k] * kpG[i,j,k].n[3]
    end
    return copy(Hout)
end

function zcross_t2c(Hin::AbstractArray{ComplexF64,4},kpG::AbstractArray{KVec,3})::AbstractArray{ComplexF64,4}
    # Hout = zeros(ComplexF64,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    for i=1:Nx,j=1:Ny,k=1:Nz
        @inbounds Hout[1,i,j,k] = -Hin[1,i,j,k] * kpG[i,j,k].m[2] - Hin[2,i,j,k] * kpG[i,j,k].n[2]
        @inbounds Hout[2,i,j,k] =  Hin[1,i,j,k] * kpG[i,j,k].m[1] + Hin[2,i,j,k] * kpG[i,j,k].n[1]
        @inbounds Hout[3,i,j,k] = 0.0
    end
    return copy(Hout)
end

function kcross_t2c(Hin::AbstractArray{ComplexF64,4},kpG::AbstractArray{KVec,3})::AbstractArray{ComplexF64,4}
    # Hout = Array{ComplexF64}(undef,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    for i=1:Nx,j=1:Ny,k=1:Nz
        @inbounds scale = -kpG[i,j,k].mag
        @inbounds Hout[1,i,j,k] = ( Hin[1,i,j,k] * kpG[i,j,k].n[1] - Hin[2,i,j,k] * kpG[i,j,k].m[1] ) * scale
        @inbounds Hout[2,i,j,k] = ( Hin[1,i,j,k] * kpG[i,j,k].n[2] - Hin[2,i,j,k] * kpG[i,j,k].m[2] ) * scale
        @inbounds Hout[3,i,j,k] = ( Hin[1,i,j,k] * kpG[i,j,k].n[3] - Hin[2,i,j,k] * kpG[i,j,k].m[3] ) * scale
    end
    return copy(Hout)
end

function kcross_c2t(Hin::AbstractArray{ComplexF64,4},kpG::AbstractArray{KVec,3})::AbstractArray{ComplexF64,4}
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
    for i=1:Nx,j=1:Ny,k=1:Nz
        @inbounds scale = kpG[i,j,k].mag
        @inbounds at1 = Hin[1,i,j,k] * kpG[i,j,k].m[1] + Hin[2,i,j,k] * kpG[i,j,k].m[2] + Hin[3,i,j,k] * kpG[i,j,k].m[3]
        @inbounds at2 = Hin[1,i,j,k] * kpG[i,j,k].n[1] + Hin[2,i,j,k] * kpG[i,j,k].n[2] + Hin[3,i,j,k] * kpG[i,j,k].n[3]
        @inbounds Hout[1,i,j,k] =  -at2 * scale
        @inbounds Hout[2,i,j,k] =  at1 * scale
    end
    return copy(Hout)
end

function kcrossinv_t2c(Hin::AbstractArray{ComplexF64,4},kpG::AbstractArray{KVec,3})::AbstractArray{ComplexF64,4}
    # Hout = Array{ComplexF64}(undef,(3,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    for i=1:Nx,j=1:Ny,k=1:Nz
        @inbounds scale = 1 / kpG[i,j,k].mag
        @inbounds Hout[1,i,j,k] = ( Hin[1,i,j,k] * kpG[i,j,k].n[1] - Hin[2,i,j,k] * kpG[i,j,k].m[1] ) * scale
        @inbounds Hout[2,i,j,k] = ( Hin[1,i,j,k] * kpG[i,j,k].n[2] - Hin[2,i,j,k] * kpG[i,j,k].m[2] ) * scale
        @inbounds Hout[3,i,j,k] = ( Hin[1,i,j,k] * kpG[i,j,k].n[3] - Hin[2,i,j,k] * kpG[i,j,k].m[3] ) * scale
    end
    return copy(Hout)
end

function kcrossinv_c2t(Hin::AbstractArray{ComplexF64,4},kpG::AbstractArray{KVec,3})::AbstractArray{ComplexF64,4}
    # Hout = Array{ComplexF64}(undef,(2,Nx,Ny,Nz))
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,2,Nx,Ny,Nz)
    for i=1:Nx,j=1:Ny,k=1:Nz
        @inbounds scale = -1 / kpG[i,j,k].mag
        @inbounds at1 = Hin[1,i,j,k] * kpG[i,j,k].m[1] + Hin[2,i,j,k] * kpG[i,j,k].m[2] + Hin[3,i,j,k] * kpG[i,j,k].m[3]
        @inbounds at2 = Hin[1,i,j,k] * kpG[i,j,k].n[1] + Hin[2,i,j,k] * kpG[i,j,k].n[2] + Hin[3,i,j,k] * kpG[i,j,k].n[3]
        @inbounds Hout[1,i,j,k] =  -at2 * scale
        @inbounds Hout[2,i,j,k] =  at1 * scale
    end
    return copy(Hout)
end

function ε⁻¹_dot(Hin::AbstractArray{ComplexF64,4},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3})::AbstractArray{ComplexF64,4}
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    for i=1:Nx,j=1:Ny,k=1:Nz
        @inbounds Hout[1,i,j,k] =  ε⁻¹[i,j,k][1,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][2,1]*Hin[2,i,j,k] + ε⁻¹[i,j,k][3,1]*Hin[3,i,j,k]
        @inbounds Hout[2,i,j,k] =  ε⁻¹[i,j,k][1,2]*Hin[1,i,j,k] + ε⁻¹[i,j,k][2,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][3,2]*Hin[3,i,j,k]
        @inbounds Hout[3,i,j,k] =  ε⁻¹[i,j,k][1,3]*Hin[1,i,j,k] + ε⁻¹[i,j,k][2,3]*Hin[2,i,j,k] + ε⁻¹[i,j,k][3,3]*Hin[3,i,j,k]
        # @inbounds Hout[1,i,j,k] =  ε⁻¹[i,j,k][1,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][1,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][1,3]*Hin[3,i,j,k]
        # @inbounds Hout[2,i,j,k] =  ε⁻¹[i,j,k][2,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][2,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][2,3]*Hin[3,i,j,k]
        # @inbounds Hout[3,i,j,k] =  ε⁻¹[i,j,k][3,1]*Hin[1,i,j,k] + ε⁻¹[i,j,k][3,2]*Hin[2,i,j,k] + ε⁻¹[i,j,k][3,3]*Hin[3,i,j,k]
    end
    return copy(Hout)
end

function ε_dot_approx(Hin::AbstractArray{ComplexF64,4},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3})::AbstractArray{ComplexF64,4}
    # Hout = similar(Hin)
    Nx,Ny,Nz = size(Hin)[2:4]
    Hout = Zygote.Buffer(Hin,3,Nx,Ny,Nz)
    for i=1:Nx,j=1:Ny,k=1:Nz
        @inbounds ε_ave = 3 / tr(ε⁻¹[i,j,k])
        @inbounds Hout[1,i,j,k] =  ε_ave * Hin[1,i,j,k]
        @inbounds Hout[2,i,j,k] =  ε_ave * Hin[2,i,j,k]
        @inbounds Hout[3,i,j,k] =  ε_ave * Hin[3,i,j,k]
    end
    return copy(Hout)
end

function M(Hin::AbstractArray{ComplexF64,4},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3};ds::MaxwellData)::Array{ComplexF64,4}
    d = ds.𝓕 * kcross_t2c(Hin,ds.kpG);
    e = ε⁻¹_dot(d,ε⁻¹); # (-1/(π)) .*
    kcross_c2t(ds.𝓕⁻¹ * e,ds.kpG)
end

function M(Hin::AbstractArray{ComplexF64,1},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3};ds::MaxwellData)::Array{ComplexF64,1}
    HinA = reshape(Hin,(2,ds.Nx,ds.Ny,ds.Nz))
    HoutA = M(HinA,ε⁻¹;ds)
    return vec(HoutA)
end

M̂(ε⁻¹::Array{SHermitianCompact{3,Float64,6},3};ds::MaxwellData) = LinearMap{ComplexF64}(H::AbstractArray{ComplexF64,1} -> M(H,ε⁻¹;ds)::AbstractArray{ComplexF64,1},(2*ds.Nx*ds.Ny*ds.Nz),ishermitian=true,ismutating=false)

function P(Hin::AbstractArray{ComplexF64,4},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ds::MaxwellData)::Array{ComplexF64,4}
    e = ds.𝓕⁻¹ * kcrossinv_t2c(Hin,ds);
    d = ε_dot_approx(e,ε⁻¹,ds); # (-1/(π)) .*
    kcrossinv_c2t(ds.𝓕 * d,ds)
end

function P(Hin::AbstractArray{ComplexF64,1},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ds::MaxwellData)::Array{ComplexF64,1}
    HinA = reshape(Hin,(2,ds.Nx,ds.Ny,ds.Nz))
    HoutA = P(HinA,ε⁻¹,ds)
    return vec(HoutA)
end

P̂(ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ds::MaxwellData) = LinearMap{ComplexF64}(H::AbstractArray{ComplexF64,1} -> P(H,ε⁻¹,ds)::AbstractArray{ComplexF64,1},(2*ds.Nx*ds.Ny*ds.Nz),ishermitian=true,ismutating=false)

function Mₖ(Hin::AbstractArray{ComplexF64,4},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},kpG::AbstractArray{KVec,3})::Array{ComplexF64,4}
    d = fft(zcross_t2c(Hin,kpG),(2:4));
    # e = ε⁻¹_dot(d,ε⁻¹);
    # kcross_c2t(ifft(e,(2:4)),kpG)
    e = ifft(ε⁻¹_dot(d,ε⁻¹),(2:4));
    kcross_c2t(e,kpG)
end

function Mₖ(Hin::AbstractArray{ComplexF64,4},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},k::Float64,g::MaxwellGrid)::Array{ComplexF64,4}
    Mₖ(Hin::AbstractArray{ComplexF64,4},ε⁻¹,kpG(SVector(0.,0.,k),Zygote.dropgrad(g)))
end

function Mₖ(Hin::AbstractArray{ComplexF64,4},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ds::MaxwellData)::Array{ComplexF64,4}
    Mₖ(Hin::AbstractArray{ComplexF64,4},ε⁻¹,ds.kpG)
end

function Mₖ(Hin::AbstractArray{ComplexF64,1},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},kpG::AbstractArray{KVec,3})::Array{ComplexF64,1}
    Nx,Ny,Nz = size(ε⁻¹)
    HinA = reshape(Hin,(2,Nx,Ny,Nz))
    HoutA = Mₖ(HinA,ε⁻¹,kpG)
    return -vec(HoutA)
end

function Mₖ(Hin::AbstractArray{ComplexF64,1},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},k::Float64,g::MaxwellGrid)::Array{ComplexF64,1}
    Nx,Ny,Nz = size(ε⁻¹)
    HinA = reshape(Hin,(2,Nx,Ny,Nz))
    return -vec(Mₖ(HinA,ε⁻¹,kpG(SVector(0.,0.,k),Zygote.dropgrad(g))))
end

function Mₖ(Hin::AbstractArray{ComplexF64,1},ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ds::MaxwellData)::Array{ComplexF64,1}
    Mₖ(Hin::AbstractArray{ComplexF64,1},ε⁻¹,ds.kpG)
end

M̂ₖ(ε⁻¹::Array{SHermitianCompact{3,Float64,6},3},ds::MaxwellData) = LinearMap{ComplexF64}(H::AbstractArray{ComplexF64,1} -> Mₖ(H,ε⁻¹,ds)::AbstractArray{ComplexF64,1},(2*ds.Nx*ds.Ny*ds.Nz),ishermitian=true,ismutating=false)

#########################
#
#   old stuff
#
#########################

# """
#     t2c: v (transverse vector) → a (cartesian vector)
# """
# function t2c(v::SVector{3,ComplexF64},k::KVec)::SVector{3,ComplexF64}
#     return v[1] * k.m + v[2] * k.n
# end

# """
#     c2t: a (cartesian vector) → v (transverse vector)
# """
# function c2t(a::SVector{3,ComplexF64},k::KVec)::SVector{3,ComplexF64}
#     v0 = a ⋅ k.m
#     v1 = a ⋅ k.n
#     return SVector(v0,v1)
# end

# """
#     kcross_t2c: a (cartesian vector) = k × v (transverse vector)
# """
# function kcross_t2c(v::SVector{3,ComplexF64},k::KVec)::SVector{3,ComplexF64}
#     return ( v[1] * k.n - v[2] * k.m ) * k.mag
# end

# """
#     kcross_t2c!: a (cartesian vector) = k × v (transverse vector)
# """
# function kcross_t2c!(v::SVector{3,ComplexF64},k::KVec)::SVector{3,ComplexF64}
#     return v = ( v[1] * k.n - v[2] * k.m ) * k.mag
# end

# """
#     kcross_c2t: v (transverse vector) = k × a (cartesian vector)
# """
# function kcross_c2t(a::SVector{3,ComplexF64},k::KVec)::SVector{3,ComplexF64}
#     at1 = a ⋅ k.m
#     at2 = a ⋅ k.n
#     v0 = -at2 * k.mag
#     v1 = at1 * k.mag
#     return SVector(v0,v1,0.0)
# end

# """
#     kcrossinv_t2c: compute a⃗ (cartestion vector) st. v⃗ (cartesian vector from two trans. vector components) ≈ k⃗ × a⃗
#     This neglects the component of a⃗ parallel to k⃗ (not available by inverting this cross product)
# """
# function kcrossinv_t2c(v::SVector{3,ComplexF64},k::KVec)::SVector{3,ComplexF64}
#     return ( v[1] * k.n - v[2] * k.m ) * ( -1 / k.mag )
# end

# """
#     kcrossinv_c2t: compute  v⃗ (transverse 2-vector) st. a⃗ (cartestion 3-vector) = k⃗ × v⃗
#     This cross product inversion is exact because v⃗ is transverse (perp.) to k⃗
# """
# function kcrossinv_c2t(a::SVector{3,ComplexF64},k::KVec)::SVector{3,ComplexF64}
#     at1 = a ⋅ k.m
#     at2 = a ⋅ k.n
#     v0 = -at2 * (-1 / k.mag )
#     v1 = at1 * ( -1 / k.mag )
#     return SVector(v0,v1,0.0)
# end

# """
#     ucross_t2c: a (cartesian vector) = u × v (transverse vector)
# """
# function ucross_t2c(u::SVector{3,ComplexF64},v::SVector{3,ComplexF64},k::KVec)::SVector{3,ComplexF64}
#     return cross(u,t2c(v,k))
# end

# """
#     d_from_H(H,k): cartesian position space d vector field from transverse, PW basis H vector field
# """
# function d_from_H(H::Array{SVector{3,ComplexF64},3},ds::MaxwellData)::Array{SVector{3,ComplexF64},3}
#     # d_recip = [ kcross_t2c(H[i,j],ds.kpG[i,j]) for i=1:ds.Nx, j=1:ds.Ny]
#     # temp =  (-1/(2π)) .* fft( reinterpret( ComplexF64, reshape( d_recip , (1,ds.Nx,ds.Ny) )), (2,3))
#     # return reshape(reinterpret(SVector{3,ComplexF64},temp),(ds.Nx,ds.Ny))
#     return reshape(reinterpret(SVector{3,ComplexF64}, (-1/(2π)) .* ( ds.𝓕 * reinterpret( ComplexF64, reshape( kcross_t2c.(H,ds.kpG), (1,ds.Nx,ds.Ny) )) ) ),(ds.Nx,ds.Ny))
# end

# """
#     e_from_d(d,ε⁻¹): e-field from d-field in cartesian position space, from division by local ε tensor
# """
# function e_from_d(d::Array{SVector{3,ComplexF64},3},ε⁻¹::Array{SHM3,2})::Array{SVector{3,ComplexF64},3}
#     return ε⁻¹ .* d
# end

# """
#     H_from_e(e,k): reciprocal space (transverse basis) H vector field from position space cartesian basis e vector field
# """
# function H_from_e(e::Array{SVector{3,ComplexF64},3},ds::MaxwellData)::Array{SVector{3,ComplexF64},3}
#     # temp = (1/(2π)) .* conj.(ifft(reinterpret(ComplexF64,reshape(e,(1,ds.Nx,ds.Ny))), (2,3)))
#     # e_recip = reshape(reinterpret(SVector{3,ComplexF64},temp),(ds.Nx,ds.Ny))
#     # return [ kcross_c2t(e_recip[i,j],ds.kpG[i,j]) for i=1:ds.Nx, j=1:ds.Ny]
#     return kcross_c2t.( reshape(reinterpret(SVector{3,ComplexF64}, (1/(2π)) .* conj.( ds.𝓕⁻¹ * reinterpret(ComplexF64,reshape(e,(1,ds.Nx,ds.Ny))) ) ) ,(ds.Nx,ds.Ny)) , ds.kpG)
# end

# """
#     d_from_e(d,ε): e-field from d-field in cartesian position space, from division by local ε tensor
# """
# function d_from_e(d::Array{SVector{3,ComplexF64},3},ε::Array{SHM3,2})::Array{SVector{3,ComplexF64},3}
#     return ε .* d
# end

# """
#     h_from_H(e,k):  position space cartesian basis h vector field from reciprocal space (transverse basis) H vector field
# """
# function h_from_H(H,ds::MaxwellData)
#     h_recip = [ t2c(H[i,j],ds.kpG[i,j]) for i=1:ds.Nx, j=1:ds.Ny]
#     temp =   fft( reinterpret( ComplexF64, reshape( h_recip , (1,ds.Nx,ds.Ny) )), (2,3))
#     return reshape(reinterpret(SVector{3,ComplexF64},temp),(ds.Nx,ds.Ny))
# end

# function flatten(H::Array{SVector{3,ComplexF64},3})
#     # return reinterpret(ComplexF64,vec(H))
#     return reinterpret(ComplexF64,vec(permutedims(H,[2,1])))
# end

# function unflatten(Hvec,ds::MaxwellData)::Array{SVector{3,ComplexF64},3}
#     # return reshape(reinterpret(SVector{3,ComplexF64},Hvec),(ds.Nx,ds.Ny))
#     return permutedims(reshape(reinterpret(SVector{3,ComplexF64},Hvec),(ds.Ny,ds.Nx)), [2,1])
# end

# Define Maxwell operator function and LinearMap instantiator objects

# function M!(H::Array{SVector{3,ComplexF64}},ε⁻¹::Array{SHM3,2},ds::MaxwellData)::Array{SVector{3,ComplexF64}}
#     H .= kcross_t2c.(H,k);
#     ds.𝓕 * H;
#     H .= ε⁻¹ .* H;
#     ds.𝓕⁻¹ * H;
#     H .= kcross_c2t.(H,k)
# end

# function M!(Hv::Vector{SVector{3,ComplexF64}},ε⁻¹::Array{SHM3,2},ds::MaxwellData)::Vector{SVector{3,ComplexF64}}
#     H = reshape(Hv,size(k))
#     H .= kcross_t2c.(H,k);
#     ds.𝓕 * H;
#     H .= ε⁻¹ .* H;
#     ds.𝓕⁻¹ * H;
#     H .= kcross_c2t.(H,k)
#     return Hv
# end

# function M!(Hv::Vector{ComplexF64},ε⁻¹::Array{SHM3,2},ds::MaxwellData)::Vector{ComplexF64}
#     HSv = copy(reinterpret(SVector{3,ComplexF64}, Hv))
#     M!(HSv,ε⁻¹,ds);
#     Hv .= copy(reinterpret(ComplexF64, HSv))
# end

# function M!(Hv::Vector{ComplexF64},ε⁻¹::Array{SHM3,2},ds::MaxwellData,Hw::Array{SVector{3,ComplexF64}})::Vector{ComplexF64}
#     copyto!(Hw, reinterpret(SVector{3,ComplexF64},reshape(Hv,(3*size(k)[1],size(k)[2:end]...))) )
#     M!(Hw,ε⁻¹,ds);
#     copyto!(Hv, vec( reinterpret(ComplexF64,Hw) ) )
#     return Hv
# end

# function M̂ₖ!(ε⁻¹::Array{SHM3,3},ds::MaxwellData,Hw::Array{SVector{3,ComplexF64}})::LinearMaps.FunctionMap{ComplexF64}
#     N = 3 * ds.Nx * ds.Ny * ds.Nz
#     f = H::Vector{ComplexF64} -> M!(H,ε⁻¹,ds,Hw)
#     return LinearMap{ComplexF64}(f,f,N,N,ishermitian=true,ismutating=true)
# end

# function M̂ₖ(ε⁻¹::Array{SHM3,2},ds::MaxwellData)::LinearMaps.FunctionMap{ComplexF64}
#     N = 2 * ds.Nx * ds.Ny * ds.Nz
#     f = H -> flatten(H_from_e( e_from_d( d_from_H( unflatten(H,ds), ds ), ε⁻¹ ), ds))
#     fc = H -> flatten(H_from_e( d_from_e( d_from_H( unflatten(H,ds), ds ), SHM3.(inv.(ε⁻¹)) ), ds))
#     return LinearMap{ComplexF64}(f,fc,N,N,ishermitian=true)
# end

# function z_cross_H(H::Array{SVector{3,ComplexF64},3},ds::MaxwellData)::Array{SVector{3,ComplexF64},3}
#     z_cross_h_recip = [ ucross_t2c(SVector{3,ComplexF64}(0.,0.,1.),H[i,j,k],ds.kpG[i,j,k]) for i=1:ds.Nx, j=1:ds.Ny, k=1:ds.Nz]
#     temp =  (-1/(2π)) .* fft( reinterpret( ComplexF64, reshape( z_cross_h_recip , (1,ds.Nx,ds.Ny,ds.Nz) )), (2,3))
#     return reshape(reinterpret(SVector{3,ComplexF64},temp),(ds.Nx,ds.Ny,ds.Nz))
# end

# function ∂ₖM̂ₖ(ε⁻¹::Array{SHM3,3},ds::MaxwellData)::LinearMaps.FunctionMap{ComplexF64}
#     N = 2 * ds.Nx * ds.Ny * ds.Nz
#     f = H -> flatten(H_from_e( e_from_d( z_cross_H( unflatten(H,ds), ds ), ε⁻¹ ), ds))
#     return LinearMap{ComplexF64}(f,f,N,N,ishermitian=true)
# end

# Define approximate inversion operator P̂ₖ ≈ M̂ₖ⁻¹ to use as preconditioner

# function P!(H::Array{SVector{3,ComplexF64}},ε::Array{SHM3,3},ds::MaxwellData)::Array{SVector{3,ComplexF64}}
#     H .= kcrossinv_t2c.(H,k);
#     ds.𝓕⁻¹ * H;
#     H .= ε .* H;
#     ds.𝓕 * H;
#     H .= kcrossinv_c2t.(H,k)
# end

# function P!(Hv::Vector{ComplexF64},ε::Array{SHM3,3},ds::MaxwellData,Hw::Array{SVector{3,ComplexF64}})::Vector{ComplexF64}
#     copyto!(Hw, reinterpret(SVector{3,ComplexF64},reshape(Hv,(3*size(k)[1],size(k)[2:end]...))) )
#     P!(Hw,ε,ds);
#     copyto!(Hv, vec( reinterpret(ComplexF64,Hw) ) )
#     return Hv
# end

# function P̂ₖ!(ε::Array{SHM3,3},ds::MaxwellData,Hw::Array{SVector{3,ComplexF64}})::LinearMaps.FunctionMap{ComplexF64}
#     N = 3 * ds.Nx * ds.Ny * ds.Nz
#     f = H::Vector{ComplexF64} -> P!(H,ε,ds,Hw)
#     return LinearMap{ComplexF64}(f,f,N,N,ismutating=true)
# end

# """
#     e_from_H_approx(H,k): cartesian position space d vector field from transverse, PW basis H vector field
# """
# function e_from_H_approx(H::Array{SVector{3,ComplexF64},3},ds::MaxwellData)::Array{SVector{3,ComplexF64},3}
#     # d_recip = [ kcross_t2c(H[i,j],ds.kpG[i,j]) for i=1:ds.Nx, j=1:ds.Ny]
#     # temp =  (-1/(2π)) .* fft( reinterpret( ComplexF64, reshape( d_recip , (1,ds.Nx,ds.Ny) )), (2,3))
#     # return reshape(reinterpret(SVector{3,ComplexF64},temp),(ds.Nx,ds.Ny))
#     return reshape(reinterpret(SVector{3,ComplexF64}, (-1/(2π)) .* ( ds.𝓕 * reinterpret( ComplexF64, reshape( kcrossinv_t2c.(H,ds.kpG), (1,ds.Nx,ds.Ny,ds.Nz) )) ) ),(ds.Nx,ds.Ny,ds.Nz))
# end

# """
#     H_from_d(e,k): reciprocal space (transverse basis) H vector field from position space cartesian basis e vector field
# """
# function H_from_d(e::Array{SVector{3,ComplexF64},3},ds::MaxwellData)::Array{SVector{3,ComplexF64},3}
#     # temp = (1/(2π)) .* conj.(ifft(reinterpret(ComplexF64,reshape(e,(1,ds.Nx,ds.Ny))), (2,3)))
#     # e_recip = reshape(reinterpret(SVector{3,ComplexF64},temp),(ds.Nx,ds.Ny))
#     # return [ kcross_c2t(e_recip[i,j],ds.kpG[i,j]) for i=1:ds.Nx, j=1:ds.Ny]
#     return kcrossinv_c2t.( reshape(reinterpret(SVector{3,ComplexF64}, (1/(2π)) .* conj.( ds.𝓕⁻¹ * reinterpret(ComplexF64,reshape(e,(1,ds.Nx,ds.Ny,ds.Nz))) ) ) ,(ds.Nx,ds.Ny,ds.Nz)) , ds.kpG)
# end

# """
# approximate inversion operator P̂ₖ ≈ M̂ₖ⁻¹ to use as preconditioner
# """
# function P̂ₖ(ε::Array{SHM3,3},ds::MaxwellData)::LinearMaps.FunctionMap
#     N = 2 * ds.Nx * ds.Ny * ds.Nz
#     f = H -> flatten(H_from_d( d_from_e( e_from_H_approx( unflatten(H,ds), ds ), ε ), ds))
#     fc = H -> flatten(H_from_d( e_from_d( e_from_H_approx( unflatten(H,ds), ds ), SHM3.(inv.(ε)) ), ds))
#     return LinearMap{ComplexF64}(f,fc,N,N)
# end


text:
using Rank2BinaryForms
using Test

using Nemo

@testset "Rank2BinaryForms.jl" begin

    @testset "Computing roots" begin
        i = Image_affineGamma(6)
        @test i[1] == QQBar.([-2,-1,-fmpq(1 // 2)])
        @test i[2] == QQBar.([3, -3, -1])
        @test i[3] == QQBar.([2, 4, -2])
        @test i[4] == QQBar.([fmpq(3 // 2), 2, 3])
        @test i[5] == QQBar.([fmpq(1 // 3), -1, -fmpq(1 // 3)])
        @test i[6] == QQBar.([fmpq(1 // 2), 2, -1])
        @test i[7] == QQBar.([fmpq(2 // 3), fmpq(4 // 3), 2])
        @test i[8] == QQBar.([fmpq(1 // 4), fmpq(1 // 2), -fmpq(1 // 2)])
        @test i[9] == QQBar.([fmpq(1 // 2), fmpq(3 // 4), fmpq(3 // 2)])
        @test i[10] == QQBar.([fmpq(1 // 3), fmpq(1 // 2), fmpq(2 // 3)])
    end

    @testset "Building polynomials" begin
        @testset "Symbols" begin
            i = affineGamma(6)([1,2,3])
            vars = [:x, :y];
            poly = buildpolynomial(vars)(i);
            global x = -2;
            global y = 1;
            @test eval(poly) == QQBar(0)
            global x = 1;
            global y = 5;
            @test eval(poly) == QQBar(6930)
            global x = 1;
            global y = 0;
            @test eval(poly) == QQBar(0)
            global x = 1;
            global y = -1;
            @test eval(poly) == QQBar(0)
            global x = 1;
            global y = -2;
            @test eval(poly) == QQBar(0)
            i = affineGamma(6)([1,2,4])
            poly = buildpolynomial(vars)(i);
            global x = 3;
            global y = 1;
            @test eval(poly) == QQBar(0)
            global x = -3;
            global y = 1;
            @test eval(poly) == QQBar(0)
            global x = -1;
            global y = 1;
            @test eval(poly) == QQBar(0)
        end
        @testset "Nemo fmpz" begin
            g(d) = k -> exp(2 * pi * im * k / d);
            set_getting_rootsofunity(g);
            introots(r) = Int.(round.(real(r)))
            R, vars = PolynomialRing(FlintZZ, ["x", "y"])

            r1 = introots.(affineGamma(4)([1,2,3]))
            @test first(r1) == -1
            @test buildpolynomial(vars)(r1) == vars[1]^3 * vars[2] + 2 * vars[1]^2 * vars[2]^2 + vars[1] * vars[2]^3

            r2 = introots.(affineGamma(4)([1,2,4]))
            @test first(r2) == 2
            @test buildpolynomial(vars)(r2) == vars[1]^3 * vars[2] - vars[1]^2 * vars[2]^2 - 2 * vars[1] * vars[2]^3

            r3 = introots.(affineGamma(4)([1,3,4]))
            @test first(r3) == 0
            @test buildpolynomial(vars)(r3) == vars[1]^3 * vars[2] + vars[1]^2 * vars[2]^2
        end
    end
end


text:
using Knet, LinearAlgebra, CUDA, Statistics

function Knet.KnetArray(x::CuArray{T,N}) where {T,N}
    p = Base.bitcast(Knet.Cptr, pointer(x))
    k = Knet.KnetPtr(p, sizeof(x), Int(CUDA.device().handle), x)
    KnetArray{T,N}(k, size(x))
end

aType = CUDA.functional() ? KnetArray{Float32} : Array{Float32}
# aType = CUDA.functional() ? CuArray{Float32} : Array{Float32}

function normal_init(a...; mean=0.0, std=0.1)
    w = randn(a...)
    if ndims(w) == 1
        fanin = length(w)
    elseif ndims(w) == 2
        fanin = size(w,2)
    else
        fanin = div(length(w),  a[end])
    end
    T = eltype(w)
    w .* T(std / sqrt(fanin)) .+ T(mean)
end
init_f = normal_init

mmul(w,x) = (w == 1 ? x : w == 0 ? 0 : reshape(w * reshape(x,size(x,1),:), (:, size(x)[2:end]...)))
# std2(a, μ, ϵ) = sqrt.(mean(abs2.(a .- μ), dims=1) .+ ϵ)

struct Embed; w; end

function Embed(vocabsize::Int, embedsize::Int)
    Embed(param(embedsize, vocabsize; atype=aType, init=init_f))
end

function (l::Embed)(x)
    l.w[:, x]
end

struct Linear; w; b; end

function Linear(inputsize::Int, outputsize::Int)
    Linear(param(outputsize, inputsize; atype=aType, init=init_f), param0(outputsize; atype=aType))
end

function (l::Linear)(x)
    mmul(l.w, x) .+ l.b
end

"""
    LayerNorm(size::Integer; eps=1e-5)
    LayerNorm(γ, β, ϵ)

References:
* [Ba, Kiros and Hinton 2016](https://arxiv.org/abs/1607.06450) Layer Normalization
* torch.nn.LayerNorm
* tf.keras.layers.LayerNormalization
"""
struct LayerNorm; γ; β; ϵ; end

function LayerNorm(dmodel; eps=1e-5)
    γ = param(dmodel; init=ones, atype=aType)
    β = param(dmodel; init=zeros, atype=aType)
    LayerNorm(γ, β, eps)
end

function (l::LayerNorm)(x, o...)
    μ = mean(x,dims=1)
    σ = std(x,mean=μ,dims=1)
    ϵ = eltype(x)(l.ϵ)
    l.γ .* (x .- μ) ./ (σ .+ ϵ) .+ l.β # TODO: doing x .- μ twice?
end

"""
    gelu(x)

Gaussian Error Linear Unit Activation Function
https://arxiv.org/abs/1606.08415
returns:
    x * sigm(1.702 * x)
"""
function gelu(x)
    return x * sigm(Float32(1.702) * x)
end

struct Boom
    lin1
    lin2
    dropout::Float32
    shortcut
    act
end

function Boom(inpsize::Int; dim_feedforward=2048, dropout=0.1, shortcut=false, act=gelu)
    lin1 = Linear(inpsize, dim_feedforward)

    if shortcut
        lin2 = aType(zeros(Float32, inpsize, dim_feedforward))
        for i in 1:inpsize:dim_feedforward
            lin2[ (i ÷ inpsize) + 1 , i:(i + inpsize - 1)] .= 1
        end
    else
        lin2 = Linear(dim_feedforward, inpsize)
    end

    Boom(lin1, lin2, dropout, shortcut, act)
end

function (boom::Boom)(inp)
    x = boom.act.(boom.lin1(inp))
    x = dropout(x, boom.dropout)

    if boom.shortcut
        # this part takes the output of the first linear layer and splits it into ( dim_feedforward ÷ inpsize ) parts and sums them together
        # Doing this using one matrix multiplication is cheaper than two reshapes and summation
        x = mmul(boom.lin2, x)
    else
        x = boom.lin2(x)
    end

    return x
end

struct Overparam
    linear
    nhid
end

function Overparam(nhid)
    Overparam(Linear(nhid, 2nhid), nhid)
end

function (o::Overparam)(x)
    x = o.linear(x)
    tanh.(x[1:o.nhid, :]) .* sigm.(x[o.nhid+1:end, :])
end

struct Attention
    q
    k
    v
    qs
    ks
    vs
    vq
    qln
    heads::Int
    hiddensize::Int
    depth::Int
    dropout::Real
end

function Attention(nhid; q=true, k=false, v=false, heads=1, dropout=0)
    @assert nhid % heads == 0 "Heads must divide vector evenly"

    q = q ? Linear(nhid, nhid) : false
    k = k ? Linear(nhid, nhid) : false
    v = v ? Linear(nhid, nhid) : false
    qs = param0(nhid, 1; atype=aType)
    ks = param0(nhid, 1; atype=aType)
    vs = param0(nhid, 1; atype=aType)
    vq = Overparam(nhid)
    qln = LayerNorm(nhid, eps=1e-12)
    depth = nhid ÷ heads

    Attention(q, k, v, qs, ks, vs, vq, qln, heads, nhid, depth, dropout)
end

function (a::Attention)(query, key, value; attn_mask=false)
    qs, ks, vs = sigm.(a.qs), sigm.(a.ks), sigm.(a.vs)
    vs = a.vq(vs)

    query = a.q !== false ? a.q(query) : query;
    query = a.q !== false ? a.qln(query) : query;
    key = a.k !== false ? a.k(key) : key;
    value = a.v !== false ? a.v(value) : value;

    q = qs[:] .* query
    k = ks[:] .* key
    v = vs[:] .* value

    q = permutedims(q, [3, 1, 2]) # -> B, H, T
    k = permutedims(k, [3, 1, 2]) # -> B, H, T
    v = permutedims(v, [3, 1, 2]) # -> B, H, T

    q = reshape(q, size(q, 1), a.depth, a.heads, :)
    k = reshape(k, size(k, 1), a.depth, a.heads, :)
    v = reshape(v, size(v, 1), a.depth, a.heads, :)

    mix, focus = attention(q, k, v; attn_mask=attn_mask, adropout=a.dropout) # mix -> (Tv, depth, heads, B)

    mix = reshape(mix, size(mix, 1), :, size(mix, 4))
    mix = permutedims(mix, (2, 3, 1))

    mix, focus # mix -> H, B, T
end

function attention(q, k, v; attn_mask=nothing, adropout=0.0)
    qk = bmm(q, k, transB=true)
    attn_scores = qk ./ Float32(sqrt(size(q, 1)))

    # to be checked
    if attn_mask !== false
        # attn_scores .+= attn_mask # error on backward pass, -> no method matching copyto!(::KnetArray{Float32,4}, ::Base.Broadcast.Broadcasted{Base.Broadcast.Style{AutoGrad.Value},NTuple{4,Base.OneTo{Int64}},typeof(identity),Tuple{AutoGrad.Result{KnetArray{Float32,4}}}})
        attn_scores = attn_mask .+ attn_scores
    end

    attn_scores = softmax(attn_scores, dims=2) # on K dim

    if adropout != 0
        dropout(attn_scores, adropout) # attention dropout
    end

    return bmm(attn_scores, v), attn_scores
end

struct SHARNNBlock
    attn
    ff
    dropout
    residual
    rnn
    lnstart
    lnmid
    lnmem
    lnff
    lnxff
end

function SHARNNBlock(embed_dim, hidden_dim; heads=1, dropout=0.0, residual=true, use_attn=true)

    if use_attn
        attn = Attention(embed_dim; heads=heads, dropout=dropout)
        lnmid = LayerNorm(embed_dim, eps=1e-12)
        lnmem = LayerNorm(embed_dim, eps=1e-12)
    else
        attn = nothing
        lnmid = nothing
        lnmem = nothing
    end

    ff = Boom(embed_dim; dim_feedforward=hidden_dim, dropout=dropout, shortcut=true, act=gelu)
    rnn = RNN(embed_dim, embed_dim, atype=aType)
    lnstart = LayerNorm(embed_dim, eps=1e-12)
    lnff = LayerNorm(embed_dim, eps=1e-12)
    lnxff = LayerNorm(embed_dim, eps=1e-12)

    SHARNNBlock(attn, ff, dropout, residual, rnn, lnstart, lnmid, lnmem, lnff, lnxff)
end

function (b::SHARNNBlock)(h, p_encoding, attn_mask; mem=nothing, hidden=nothing)
    h = b.lnstart(h);

    # RNN part
    if hidden !== nothing
        b.rnn.h, b.rnn.c = hidden
    else
        b.rnn.h, b.rnn.c = 0, 0
    end

    x = b.rnn(h)
    x = dropout(x, b.dropout)
    new_hidden = (value(b.rnn.h), value(b.rnn.c))
    h = b.residual ? (h + x) : x

    # Attention part

    focus, new_mem = nothing, []
    if b.attn !== nothing
        if mem !== nothing
            bigh = cat(mem, b.lnmem(h), dims=3)
        else
            bigh = b.lnmem(h)
        end

        newmemsize = length(p_encoding)
        newmem_start_idx = newmemsize < (size(bigh, 3) + 1) ? (size(bigh, 3) + 1 - newmemsize) : 1
        new_mem = value(bigh[:, :, newmem_start_idx:end])

        h = b.lnmid(h)
        x, focus = b.attn(h, bigh, bigh; attn_mask=attn_mask)
        x = dropout(x, b.dropout)
        h = x + h
    end

    # Boom layer
    h, x = b.lnff(h), b.ff(b.lnxff(h))
    x = dropout(x, b.dropout)
    h = h + x

    h, new_mem, new_hidden, focus
end

mutable struct SHARNN
    encoder
    decoder
    vocab
    blocks
    idrop
    drop
    hdrop
    num_max_positions
    pos_emb
    new_hidden
    new_mems
end

function SHARNN(embsz::Int, hidden::Int, vocab::Vocab, nlayers::Int; num_max_positions=1024, nheads=1, dropout=0.0, dropouth=0.1, dropouti=0.1, wdrop=0, tie_weights=false)

    encoder = Embed(length(vocab.i2v), embsz)

    pos_emb = zeros(num_max_positions)

    if tie_weights
        decoder = nothing
    else
        decoder = Linear(embsz, length(vocab.i2v))
    end

    blocks = []
    for i=1:nlayers
        push!(blocks, SHARNNBlock(embsz, hidden; heads=nheads, dropout=dropouth, residual=false, use_attn=(i == (nlayers - 1)))) # (i == (nlayers - 1))
    end

    SHARNN(encoder, decoder, vocab, blocks, dropouti, dropout, dropouth, num_max_positions, pos_emb, nothing, nothing)
end

function (s::SHARNN)(x; hidden=nothing, mems=nothing, return_h=true)
    embed = dropout(s.encoder(x), s.idrop)

    if mems !== nothing
        maxmem = s.num_max_positions - size(embed, 3)
        maxmem_start_idx = maxmem < (size(mems[1], 3) + 1) ? (size(mems[1], 3) + 1 - maxmem) : 1
        mems = [ m[:, :, maxmem_start_idx:end] for m in mems]
    end

    new_hidden, new_mems, focus = [], [], []
    attn_mask = fill(typemin(Float32), size(embed, 3), size(embed, 3)) # -> TxT
    attn_mask = triu(attn_mask, 1)
    if mems !== nothing
        max_mem_size = max([size(m, 3) for m in mems]...)
        attn_mask = cat(zeros(Float32, size(embed, 3), max_mem_size), attn_mask, dims=2)
    end
    attn_mask = aType(attn_mask)

    h = embed
    for (i, block) in enumerate(s.blocks)
        mem = mems !== nothing ? mems[i] : nothing
        hid = hidden !== nothing ? hidden[i] : nothing
        h, m, nh, f = block(h, s.pos_emb, attn_mask; mem=mem, hidden=hid)
        push!(new_hidden, nh)
        push!(new_mems, m)
    end
    h = dropout(h, s.drop)

    h, new_hidden, new_mems
end

function (s::SHARNN)(src, tgt; average=true)
    h, new_hidden, new_mems = s(src; hidden=s.new_hidden, mems=s.new_mems)
    s.new_hidden = new_hidden
    s.new_mems = new_mems

    if s.decoder !== nothing
        scores = s.decoder(h)
    else
        scores = mmul(s.encoder.w', h)
    end

    nll(scores, tgt; dims=1, average=average)
end

nothing


text:
struct Rules{T<:Real}
	breakpoints::AbstractVector{T}
	labels::AbstractVector
	iflower::Bool
end

"""
    Rules(breakpoints::AbstractVector, labels::AbstractVector, iflower::Bool=true)

Create a container for interpretation rules of thumb. See [`interpret(x::Real, rules::Rules)`](@ref).

# Arguments
- `breakpoints`: Vector of value break points (edges defining categories).
- `labels`: Labels associated with each category. Must contain one label more than breakpoints.
- `iflower`: If true, each label will be given if the value *is lower* than its breakpoint. The contrary if false.

# Examples
```jldoctest
Rules([0.05], ["significant", "not significant"], true)

# output

Rules{Float64}([0.05], ["significant", "not significant"], true)
```
"""
function Rules(breakpoints::AbstractVector, labels::AbstractVector, iflower::Bool=true)
    if length(breakpoints) != length(labels)-1
        error("There must be exactly one more label than breakpoints.")
    elseif !issorted(breakpoints)
        error("Breakpoints must be sorted.")
    else
        Rules(breakpoints, labels, iflower)
    end
end
Rules(; breakpoints::AbstractVector, labels::AbstractVector, iflower::Bool=true) =
Rules(breakpoints, labels, iflower)

"""
    interpret(x::Real, rules::Rules)

Interpret a value based on a set of rules of thumb.

# Arguments
- `x`: The value to interpret.
- `rules`: A [`Rules`](@ref) object.

# Examples
```jldoctest
p_rules = Rules([0.05], ["significant", "not significant"], true)
interpret(0.04, p_rules)

# output

"significant"
```
"""
function interpret(x::Real, rules::Rules)
    rules.labels[1+searchsortedlast(rules.breakpoints, x, lt=(rules.iflower ? (<) : (<=)))]
end
interpret(x::Real; rules::Rules) = interpret(x, rules)


text:
## WeightVec

using ArrayViews
using StatsBase
using Base.Test

@test isa(weights([1, 2, 3]), WeightVec{Int})
@test isa(weights([1., 2., 3.]), WeightVec{Float64})

@test isempty(weights(Float64[]))

w = [1., 2., 3.]
wv = weights(w)
@test eltype(wv) == Float64
@test length(wv) == 3
@test values(wv) === w
@test sum(wv) == 6.0
@test !isempty(wv)

## wsum

x = [6., 8., 9.]
w = [2., 3., 4.]

@test wsum(Float64[], Float64[]) === 0.0
@test wsum(x, w) == 72.0

## wsum along dimensions

@test wsum(x, w, 1) == [72.0]

x = rand(6, 8)
w1 = rand(6)
w2 = rand(8)

@test size(wsum(x, w1, 1)) == (1, 8)
@test size(wsum(x, w2, 2)) == (6, 1)

@test_approx_eq wsum(x, w1, 1) sum(x .* w1, 1)
@test_approx_eq wsum(x, w2, 2) sum(x .* w2', 2)

x = rand(6, 5, 4)
w1 = rand(6)
w2 = rand(5)
w3 = rand(4)

@test size(wsum(x, w1, 1)) == (1, 5, 4)
@test size(wsum(x, w2, 2)) == (6, 1, 4)
@test size(wsum(x, w3, 3)) == (6, 5, 1)

@test_approx_eq wsum(x, w1, 1) sum(x .* w1, 1)
@test_approx_eq wsum(x, w2, 2) sum(x .* w2', 2)
@test_approx_eq wsum(x, w3, 3) sum(x .* reshape(w3, 1, 1, 4), 3)

v = view(x, 2:4, :, :)

@test_approx_eq wsum(v, w1[1:3], 1) sum(v .* w1[1:3], 1)
@test_approx_eq wsum(v, w2, 2) sum(v .* w2', 2)
@test_approx_eq wsum(v, w3, 3) sum(v .* reshape(w3, 1, 1, 4), 3)

## wsum for Arrays with non-BlasReal elements

x = rand(1:100, 6, 8)
w1 = rand(6)
w2 = rand(8)

@test_approx_eq wsum(x, w1, 1) sum(x .* w1, 1)
@test_approx_eq wsum(x, w2, 2) sum(x .* w2', 2)

## wsum!

x = rand(6)
w = rand(6)

r = ones(1)
@test wsum!(r, x, w, 1; init=true) === r
@test_approx_eq r [dot(x, w)]

r = ones(1)
@test wsum!(r, x, w, 1; init=false) === r
@test_approx_eq r [dot(x, w) + 1.0]

x = rand(6, 8)
w1 = rand(6)
w2 = rand(8)

r = ones(1, 8)
@test wsum!(r, x, w1, 1; init=true) === r
@test_approx_eq r sum(x .* w1, 1)

r = ones(1, 8)
@test wsum!(r, x, w1, 1; init=false) === r
@test_approx_eq r sum(x .* w1, 1) .+ 1.0

r = ones(6)
@test wsum!(r, x, w2, 2; init=true) === r
@test_approx_eq r sum(x .* w2', 2)

r = ones(6)
@test wsum!(r, x, w2, 2; init=false) === r
@test_approx_eq r sum(x .* w2', 2) .+ 1.0

x = rand(8, 6, 5)
w1 = rand(8)
w2 = rand(6)
w3 = rand(5)

r = ones(1, 6, 5)
@test wsum!(r, x, w1, 1; init=true) === r
@test_approx_eq r sum(x .* w1, 1)

r = ones(1, 6, 5)
@test wsum!(r, x, w1, 1; init=false) === r
@test_approx_eq r sum(x .* w1, 1) .+ 1.0

r = ones(8, 1, 5)
@test wsum!(r, x, w2, 2; init=true) === r
@test_approx_eq r sum(x .* w2', 2)

r = ones(8, 1, 5)
@test wsum!(r, x, w2, 2; init=false) === r
@test_approx_eq r sum(x .* w2', 2) .+ 1.0

r = ones(8, 6)
@test wsum!(r, x, w3, 3; init=true) === r
@test_approx_eq r sum(x .* reshape(w3, (1, 1, 5)), 3)

r = ones(8, 6)
@test wsum!(r, x, w3, 3; init=false) === r
@test_approx_eq r sum(x .* reshape(w3, (1, 1, 5)), 3) .+ 1.0

## the sum and mean syntax

@test_approx_eq sum([1.0, 2.0, 3.0], weights([1.0, 0.5, 0.5])) 3.5
@test_approx_eq sum(1:3, weights([1.0, 1.0, 0.5])) 4.5

@test_approx_eq mean([1:3], weights([1.0, 1.0, 0.5])) 1.8
@test_approx_eq mean(1:3, weights([1.0, 1.0, 0.5])) 1.8

a = reshape(1.0:27.0, 3, 3, 3)
for wt in ([1.0, 1.0, 1.0], [1.0, 0.2, 0.0], [0.2, 0.0, 1.0])
	@test_approx_eq sum(a, weights(wt), 1) sum(a.*reshape(wt, length(wt), 1, 1), 1)
	@test_approx_eq sum(a, weights(wt), 2) sum(a.*reshape(wt, 1, length(wt), 1), 2)
	@test_approx_eq sum(a, weights(wt), 3) sum(a.*reshape(wt, 1, 1, length(wt)), 3)
	@test_approx_eq mean(a, weights(wt), 1) sum(a.*reshape(wt, length(wt), 1, 1), 1)/sum(wt)
	@test_approx_eq mean(a, weights(wt), 2) sum(a.*reshape(wt, 1, length(wt), 1), 2)/sum(wt)
	@test_approx_eq mean(a, weights(wt), 3) sum(a.*reshape(wt, 1, 1, length(wt)), 3)/sum(wt)
	@test_throws ErrorException mean(a, weights(wt), 4)
end