"""
    get_quantile_scores(scores, quantiles = 0.0:0.01:1.0)

return the quantiles of the given N dimensional anomaly `scores` cube. `quantiles` (default: `quantiles = 0.0:0.01:1.0`) is a Float range of quantiles.
Any score being greater or equal `quantiles[i]` and beeing smaller than `quantiles[i+1]` is assigned to the respective quantile `quantiles[i]`.

# Examples
```
julia> scores1 = rand(10, 2)
julia> quantile_scores1 = get_quantile_scores(scores1)
```
"""
function get_quantile_scores(scores::AbstractArray{tp,N}, quantiles::StepRangeLen{Float64} = 0.0:0.01:1.0) where {tp,N}
  quantile_scores = zeros(Float64, size(scores))
  get_quantile_scores!(quantile_scores, scores, quantiles)
  return quantile_scores
end

"""
    get_quantile_scores!{tp,N}(quantile_scores::AbstractArray{Float64, N}, scores::AbstractArray{tp,N}, quantiles::StepRangeLen{Float64} = 0.0:0.01:1.0)

return the quantiles of the given N dimensional `scores` array into a preallocated `quantile_scores` array, see `get_quantile_scores()`.
"""
function get_quantile_scores!(quantile_scores::AbstractArray{Float64, N}, scores::AbstractArray{tp,N}, quantiles::StepRangeLen{Float64} = 0.0:0.01:1.0; return_thresholds::Bool = false) where {tp,N}
  LENGTH = length(scores)
  thresholds = quantile(reshape(scores, prod(size(scores))), collect(quantiles))
  n_quants = size(thresholds, 1)
  for j = 1:LENGTH
      if(scores[j] <= thresholds[1])
        quantile_scores[j] = quantiles[1]
      end
      for i = 2:n_quants
        if(scores[j] > thresholds[i-1] && scores[j] <= thresholds[i])
          quantile_scores[j] = quantiles[i]
        end
      end
  end
  if return_thresholds == false
    return quantile_scores
  else
    return thresholds
  end
end

"""
    compute_ensemble(m1_scores, m2_scores[, m3_scores, m4_scores]; ensemble = "mean")

compute the mean (`ensemble = "mean"`), minimum (`ensemble = "min"`), maximum (`ensemble = "max"`) or median (`ensemble = "median"`) of the given anomaly scores.
Supports between 2 and 4 scores input arrays (`m1_scores, ..., m4_scores`). The scores of the different anomaly detection algorithms should be somehow comparable,
e.g., by using `get_quantile_scores()` before.

# Examples
```
julia> using MultivariateAnomalies
julia> scores1 = rand(10, 2)
julia> scores2 = rand(10, 2)
julia> quantile_scores1 = get_quantile_scores(scores1)
julia> quantile_scores2 = get_quantile_scores(scores2)
julia> compute_ensemble(quantile_scores1, quantile_scores2, ensemble = "max")
```
"""
function compute_ensemble(m1_scores::Array{Tp, N}, m2_scores::Array{Tp, N}; ensemble = "mean") where {Tp, N}
    @assert any(ensemble .== ["mean","min","max","median"])

    scores = cat(m1_scores, m2_scores, dims = N+1);

    if ensemble == "mean"
      ensemble_scores = dropdims(mean(scores, dims = N+1), dims = N+1)
    end

    if ensemble == "median"
      ensemble_scores = dropdims(median(scores, dims = N+1), dims = N+1)
    end

    if ensemble == "max"
      ensemble_scores = dropdims(maximum(scores, dims = N+1), dims = N+1)
    end

    if ensemble == "min"
      ensemble_scores = dropdims(minimum(scores, dims = N+1), dims = N+1)
    end
  return ensemble_scores
end

function compute_ensemble(m1_scores::Array{Tp, N}, m2_scores::Array{Tp, N}, m3_scores::Array{Tp, N}; ensemble = "mean") where {Tp, N}
    @assert any(ensemble .== ["mean","min","max","median"])

    scores = cat(m1_scores, m2_scores, m3_scores, dims = N+1);

    if ensemble == "mean"
      ensemble_scores = dropdims(mean(scores, dims = N+1), dims = N+1)
    end

    if ensemble == "median"
      ensemble_scores = dropdims(median(scores, dims = N+1), dims = N+1)
    end

    if ensemble == "max"
      ensemble_scores = dropdims(maximum(scores, dims = N+1), dims = N+1)
    end

    if ensemble == "min"
      ensemble_scores = dropdims(minimum(scores, dims = N+1), dims = N+1)
    end
  return ensemble_scores
end

function compute_ensemble(m1_scores::Array{Tp, N}, m2_scores::Array{Tp, N}, m3_scores::Array{Tp, N}, m4_scores::Array{Tp, N}; ensemble = "mean") where {Tp, N}
    @assert any(ensemble .== ["mean","min","max","median"])

    scores = cat(m1_scores, m2_scores, m3_scores, m4_scores, dims = N+1);

    if ensemble == "mean"
      ensemble_scores = dropdims(mean(scores, dims = N+1), dims = N+1)
    end

    if ensemble == "median"
      ensemble_scores = dropdims(median(scores, dims = N+1), dims = N+1)
    end

    if ensemble == "max"
      ensemble_scores = dropdims(maximum(scores, dims = N+1), dims = N+1)
    end

    if ensemble == "min"
      ensemble_scores = dropdims(minimum(scores, dims = N+1), dims = N+1)
    end
  return ensemble_scores
end

###################################
# end


module Disease
    export epidemic

    function epidemic(tm::Number, n::Number, s0::Number, i0::Number, b::Number, a::Number)
        dt::Float64 = tm/n
        dta::Float64 = dt * a
        dtb::Float64 = -dt * b
        s::Float64 = s0
        i::Float64 = i0
        r::Float64 = 0
        for t::Float64 = 0:dt:tm
            ds = dtb * s * i
            dr = dta * i
            di = -ds -dr
            if di < 0
                break
            end
            s += ds
            r += dr
            i += di
        end
        return i
    end
end


function [ tb,qb,qbs,tbv,pisfc, pib, rhou, rhow,rhb, ub,um,u ] = ...
    base( profile_method,nx,nz,dz,psurf,qsurf,q4km,ztr,temptr,ttr,tsurf,p0,cp,g,...
    rd,xk,c_v,zu,rhow,rhou,tb,tbv,qb,qbs,rhb,pib,ub,um,u)
#UNTITLED5 Summary of this function goes here
#   Detailed explanation goes here

if profile_method == 1
    #########################################
    # Define a dry and neutral environment

    # Set base state theta to 300 K and mixing ratio to 0 kg/kg

    for k = 2 : nz - 1
        tb[k] = 300.0;
        qb[k] = 0.0;
        tbv[k] = tb(k)*(1 + 0.61*qb(k));

    end

    # Assign pi
    pisfc = (psurf / p0)^xk;
    pib[2] = pisfc - g*0.5*dz/(cp*tbv(2));

    for k = 3 : nz - 1

        tbvavg[k] = 0.5*(tbv(k) + tbv(k-1));
        pib[k] = pib(k-1) - g*dz/(cp*tbvavg(k));

    end

    for k = 2 : nz - 1
        #  Assign mean state density at u point

        rhou[k] = (p0*pib(k)^(c_v/rd))/(rd*tbv(k));

        # Assign mean state density at w point

        rhow[k] = 0.5*(rhou(k) + rhou(k-1));

    end
    #Ensure there is no gradient at the boundaries
    tb[1] = tb(2);
    tb[nz] = tb(nz-1);
    pib[1] = pib(2);
    pib[nz] = pib(nz-1);
    rhou[1] = rhou(2);
    rhou[nz] = rhou(nz-1);
    rhow[2] = 1.15;
    rhow[1] = rhow(2);
    rhow[nz] = rhow(nz-1);

elseif profile_method == 2
    ## fovellwk82soundingdef
    # define a sounding on the u/scalar points

    # Assign Values to vertical grid

    for k = 2 : nz - 1

        # Assign water vapor mixing ratio
        if(zu[k] <= 4000.0)
            qb[k] = qsurf - (qsurf - q4km)*zu(k)/4000.0;
        elseif (zu[k] <= 8000.0)
            qb[k] = q4km - q4km*(zu(k)-4000.0)/4000.0;
        else
            qb[k] = 0.0;
        end

        # Assign potential temperature
        if(zu[k]<= ztr)
            tb[k] = tsurf + (ttr - tsurf)*(zu(k)/ztr)^1.25;
        else
            tb[k] = ttr*exp(g*(zu(k) - ztr)/(cp*temptr));
        end

        # Assign virtual potential temperature
        tbv[k] = tb(k)*(1 + 0.61*qb(k));

    end

    # Assign pi
    pisfc = (psurf / p0)^xk;
    pib[2] = pisfc - g*0.5*dz/(cp*tbv(2));

    for k = 3 : nz - 1

        tbvavg[k] = 0.5*(tbv(k) + tbv(k-1));
        pib[k] = pib(k-1) - g*dz/(cp*tbvavg(k));

    end

    #  Assign mean state density at u point
    for k = 2 : nz-1

        rhou[k] = (p0*pib(k)^(c_v/rd))/(rd*tbv(k));

        # Assign mean water vapor saturation mixing ratio

        t[k] = tb(k)*pib(k);
        p[k] = p0*pib(k)^(cp/rd);
        qbs[k] = (380.0/p(k))*exp(17.27*(t(k)-273.0)/(t(k)-36.0));

        # Calculate and Assign RH

        rhb[k] = qb(k)/qbs(k);

    end

    #  define density at the true surface from known information
    pisfc = (psurf/p0)^xk;
    rhow[2] = p0*(pisfc^(c_v/rd))/(rd*tsurf);

    #  define density at the other w points by interpolating from u/scalar points
    for k=3:nz-1
        rhow[k] = 0.5*(rhou(k) + rhou(k-1));
    end

elseif profile_method == 3

        #  define the Weisman Klemp sounding on the u/scalar points

        for k=2 : nz-1
            #  assign potential temperatures and relative humidities
            if zu[k]<= ztr
                tb[k] = tsurf + (ttr-tsurf)*(zu[k]/ztr)^1.25;
                rhb[k] = 1 - 0.02*(zu[k]/ztr)^1.25;
            elseif zu[k] > ztr
                tb[k] = ttr*exp(g*(zu[k]-ztr)/(cp*temptr));
                rhb[k] = 0.98;
            end
            #  first guess sans moisture
            tbv[k] = tb[k];
        end

        #  compute the non-dimensional pressure profile as a first guess
        #    (needed for qvs below)
        pisfc = (psurf / p0)^xk;
        pib[2] = pisfc - g*0.5*dz/(cp*tbv[2]);

        for k = 3 : nz - 1
            tbvavg = 0.5*(tbv[k] + tbv[k-1]);
            pib[k] = pib[k-1] - g*dz/(cp*tbvavg);
        end

        #  convert to mixing ratios
        for k=2 : nz-1
            p = p0*pib[k]^(cp/rd);
            temp = tb[k]*pib[k];
            qvs = (380.0/p) * exp((17.27*(temp-273.0))/(temp-36.0));
#             if zu(k) < 20000
            qb[k] = min(qsurf,qvs*rhb[k]);
#             else
#             qb(k) = 0;
#             end
            #  define virtual potential temperatures
            tbv[k] = tb[k]*(1 + 0.61*qb[k]);

        #   if k==98 || k == 99 || k == 100
        #     display(k)
        #     display(p)
        #     display(temp)
        #     display(qvs)
        #     display('qb')
        #     display(qb(k))
        #     display('tbv')
        #     display(tbv(k))
        #     display('rhb')
        #     display(rhb(k))
        #   end
        # end

        #  go back and recompute the non-dimensional pressure profile now that
        #    we have humidity included
        pisfc = (psurf / p0)^xk;
        pib[2] = pisfc - g*0.5*dz/(cp*tbv[2]);

        for k = 3 : nz - 1

            tbvavg = 0.5*(tbv[k] + tbv[k-1]);
            pib[k] = pib[k-1] - g*dz/(cp*tbvavg);

        end

        #  compute the density profile
        pisfc = (psurf/p0)^xk;
        rhow[2] = p0*(pisfc^(c_v/rd))/(rd*tsurf);

        #  define density at the other w points by interpolating from u/scalar points
        #  define density at u/scalar points
      for k=2:nz-1
        rhou[k] = p0*(pib[k]^(c_v/rd))/(rd*tbv[k]);
      end

#  define density at the true surface from known information
      pisfc = (psurf/p0)^xk;
      rhow[2] = p0*(pisfc^(c_v/rd))/(rd*tsurf);

#  define density at the other w points by interpolating from u/scalar points
      for k=3:nz-1
        rhow[k] = 0.5*(rhou[k] + rhou[k-1]);
      end

        #  compute the profiles of saturation mixing ratio and relative humidity
       #  define saturation mixing ratio on u/scalar points
      for k=2: nz-1
        p = p0*pib[k]^(cp/rd);
        temp = tb[k]*pib[k];
#  use Teten's formula (see Pielke's text, 2nd ed., p. 257-258 for more info)
        qbs[k] = (380.0/p) * exp((17.27*(temp-273.0))/(temp-36.0));
        rhb[k] = qb[k] / qbs[k];
      end

 end # End Method If Statement

    #Ensure there is no gradient at the boundaries
    tb[1] = tb[2];
    tb[nz] = tb[nz-1];
    tbv[1] = tbv[2];
    tbv[nz] = tbv[nz-1];
    pib[1] = pib[2];
    pib[nz] = pib[nz-1];
    rhou[1] = rhou[2];
    rhou[nz] = rhou[nz-1];
    rhow[1] = rhow[2];
    rhow[nz] = rhow[nz-1];
    qb[1] = qb[2];
    qb[nz] = qb[nz-1];

## Define U base

# Uniform Flow

    ub[:] .= -10;
    um[:,:] .= -10;
    u[:,:] .= -10;

############################################################
# High Surface Wind decreasing as altitude increases

#     usfc=-10.0;  # u-wind at the surface (m/s)
#     # * for initwinds=1 (only)
#     deltau=10.0;  # u-wind change over the layer (m/s)
#     # * for initwinds=1 (only)
#     shrdepth=3000.0;  # depth of deltau layer (m)
#     #  * for initwinds=1 (only)
#
#
#     for k=2:nz-1
#         #  assign base state values
#         if zu(k)<shrdepth
#             ub(k) = usfc + zu(k)*deltau/shrdepth;
#         else
#             ub(k) = usfc + deltau;
#         end
#         for i= 1 :nx
#             #  copy base state values into prognosed u array
#             u(i,k) = ub(k);
#             um(i,k) = ub(k);
#         end
#     end

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

# Low-Level Jet at a given z height (given by shrdepth)
#
#     usfc=0;  # u-wind at the surface (m/s)
#     # * for initwinds=1 (only)
#     deltau=-10.0;  # u-wind change over the layer (m/s)
#     # * for initwinds=1 (only)
#     shrdepth=2000.0;  # depth of deltau layer (m)
#     #  * for initwinds=1 (only)
#
#     for k=2:nz-1
#         #  assign base state values
#         if zu(k)<=shrdepth
#             ub(k) = usfc + zu(k)*deltau/shrdepth;
#         elseif zu(k) <= 2*shrdepth
#             ub(k) = 2*deltau  - deltau*(zu(k)/(shrdepth));
#         else
#             ub(k) = 0;
#         end
#     for i= 1 :nx
#             #  copy base state values into prognosed u array
#             u(i,k) = ub(k);
#             um(i,k) = ub(k);
#      end
#      end
######################################################################
    #  handle boundaries
    for i=1:nx
        u[i,1] = ub[2];
        up[i,1] = ub[2];
        um[i,1] = ub[2];
        u[i,nz] = ub[nz-1];
        up[i,nz] = ub[nz-1];
        um[i,nz] = ub[nz-1];
    end

end


data = split(
"""
.#..##.###...#######
##.############..##.
.#.######.########.#
.###.#######.####.#.
#####.##.#.##.###.##
..#####..#.#########
####################
#.####....###.#.#.##
##.#################
#####.##.###..####..
..######..##.#######
####.##.####...##..#
.#####..#.######.###
##...#.##########...
#.##########.#######
.####.#.###.###.#.##
....##.##.###..#####
.#.#.###########.###
#.#.#.#####.####.###
###.##.####.##.#..##
""", "\n")

data = split(
"""
......#.#.
#..#.#....
..#######.
.#.#.###..
.#..#.....
..#....#.#
#..#....#.
.##.#..###
##...#..#.
.#....####
""", "\n")

function maketuples(data)
    a = Vector{Tuple{Int, Int}}()
    for (i, line) in enumerate(data)
        for j in findall(==("#"), split(line, ""))
            push!(a, (i,j))
        end
    end
    a
end

function countseen(a::Vector, i)
    seen = Set{Tuple{Float64, Float64}}()
    for asts in a
        asts == i && continue
        dir = asts .- i
        dirp = dir[1] == 0 ? dir ./ abs(dir[2]) : dir ./ abs(dir[1])
        push!(seen, dirp)
    end
    length(seen)
end

data = readlines("data/Dec10_data.txt")
a = maketuples(data)
t = countseen(a, a[1])
findmax(map(i->countseen(a, i), a))

using DataStructures
function seenDict(a::Vector, i)
    seen = MultiDict{Tuple{Float64, Float64}, Tuple{Int, Int}}()
    for asts in a
        asts == i && continue
        dir = asts .- i
        dirp = dir[1] == 0 ? dir ./ abs(dir[2]) : dir ./ abs(dir[1])
        insert!(seen, dirp, dir)
    end
    seen
end

function shootstars(myast, a)
    ret = Tuple{Int, Int}[]
    sd = seenDict(a, myast)
    for (k,v) in sd
        sort!(v, by = x->sum((x .- myast).^2), rev = true)
    end
    sortedkeys = sort(collect(keys(sd)), by = x->atan(x[2], x[1]), rev = true)
    counter = 0
    for key in sortedkeys
        if !isempty(sd[key])
            p = pop!(sd[key])
            push!(ret, p .+ myast)
            (counter += 1) == 200 && return p .+ myast .- (1,1)
        end
    end
end
res = shootstars(a[353], a)
res[2]*100 + res[1]


using OceanTurb, Printf

using OceanTurb.ModularKPP: DiagnosticPlumeModel

@use_pyplot_utils

usecmbright()

function makeplot!(axs, model)
    ModularKPP.update_state!(model)

    ΔT = model.state.plume.T - model.solution.T

    ϵ = CellField(model.grid)
    for i in eachindex(ϵ)
        @inbounds ϵ[i] = ModularKPP.entrainment(model.grid.zc[i], model)
    end

    for ax in axs
        sca(ax); cla()
    end

    sca(axs[1]); xlabel(L"\breve T - T"); ylabel(L"z \, \mathrm{(m)}")
    sca(axs[2]); xlabel(L"T")
    sca(axs[3]); xlabel(L"\breve W^2")
    sca(axs[4]); xlabel(L"\epsilon")

    sca(axs[1])
    plot(ΔT)
    removespines("top", "right")

    sca(axs[2])
    plot(model.solution.T, label="Environment \$ T \$")
    plot(model.state.plume.T, "--", label="Plume \$ \\breve T \$")
    removespines("top", "left", "right")

    sca(axs[3])
    plot(model.state.plume.W²)
    removespines("top", "left", "right")

    sca(axs[4])
    plot(ϵ)
    removespines("top", "left", "right")

    axs[2].tick_params(left=false, labelleft=false)
    axs[3].tick_params(left=false, labelleft=false)
    axs[4].tick_params(left=false, labelleft=false)

    return nothing
end

Qb = 1e-7
N² = 1e-5
Δt = 1minute

# The standard setup except with a plume model rather than a counter-gradient flux model.
model = ModularKPP.Model(;         grid = UniformGrid(N=64, H=256),
                                stepper = :BackwardEuler,
                           nonlocalflux = DiagnosticPlumeModel(Ca=0.1, Ce=0.4, Cα=1.0))

# Initial condition and fluxes
dTdz = N² / (model.constants.α * model.constants.g)
T₀(z) = 20 + dTdz*z
model.solution.T = T₀
OceanTurb.set!(model.state.plume.T, T₀)

Qθ = Qb / (model.constants.α * model.constants.g)

model.bcs.T.top = FluxBoundaryCondition(Qθ)
model.bcs.T.bottom = GradientBoundaryCondition(dTdz)

fig, axs = subplots(ncols=4, sharey=true)

run_until!(model, Δt, 24hour)
makeplot!(axs, model)


"""
    Quat{T} <: Rotation{T}
    Quat(w, x, y, z)

The `Quat` type is a 3×3 matrix representation of a normalized quaternion.
They allow you to transparently use (fast) quaternion algebra to store, compose
and invert 3D rotations, while at the same time letting you apply rotations
through matrix-vector multiplication.

Note: the constructor will always renormalize the input so that the quaternion
has length 1 (w² + x² + y² + z² = 1), and the rotation matrix is orthogonal.
"""
immutable Quat{T} <: Rotation{T}
    w::T
    x::T
    y::T
    z::T

    # For the moment we ensure that it is normalized upon construction.
    function Quat(w, x, y, z)
        norm = copysign(sqrt(w*w + x*x + y*y + z*z), w)
        # Should this be an error or warning, if it isn't approximately normalized? E.g.:
        #if norm !≈ 1
        #    error("Expected a normalized quaternion") # or warn() ?
        #end
        new(w/norm, x/norm, y/norm, z/norm)
    end
end

# StaticArrays will take over *all* the constructors and put everything in a tuple...
# but this isn't quite what we mean when we have 4 inputs (not 9).
@inline (::Type{Quat}){W,X,Y,Z}(w::W, x::X, y::Y, z::Z) = Quat{promote_type(promote_type(promote_type(W, X), Y), Z)}(w, x, y, z)

# These 2 functions are enough to satisfy the entire StaticArrays interface:
function (::Type{Q}){Q<:Quat}(t::NTuple{9})
    q = Q(sqrt(abs(1  + t[1] + t[5] + t[9])) / 2,
               copysign(sqrt(abs(1 + t[1] - t[5] - t[9]))/2, t[6] - t[8]),
               copysign(sqrt(abs(1 - t[1] + t[5] - t[9]))/2, t[7] - t[3]),
               copysign(sqrt(abs(1 - t[1] - t[5] + t[9]))/2, t[2] - t[4]))
end

function Base.getindex(q::Quat, i::Integer)
    if i == 1
        ww = (q.w * q.w)
        xx = (q.x * q.x)
        yy = (q.y * q.y)
        zz = (q.z * q.z)

        ww + xx - yy - zz
    elseif i == 2
        xy = (q.x * q.y)
        zw = (q.w * q.z)

        2 * (xy + zw)
    elseif i == 3
        xz = (q.x * q.z)
        yw = (q.y * q.w)

        2 * (xz - yw)
    elseif i == 4
        xy = (q.x * q.y)
        zw = (q.w * q.z)

        2 * (xy - zw)
    elseif i == 5
        ww = (q.w * q.w)
        xx = (q.x * q.x)
        yy = (q.y * q.y)
        zz = (q.z * q.z)

        ww - xx + yy - zz
    elseif i == 6
        yz = (q.y * q.z)
        xw = (q.w * q.x)

        2 * (yz + xw)
    elseif i == 7
        xz = (q.x * q.z)
        yw = (q.y * q.w)

        2 * (xz + yw)
    elseif i == 8
        yz = (q.y * q.z)
        xw = (q.w * q.x)

        2 * (yz - xw)
    elseif i == 9
        ww = (q.w * q.w)
        xx = (q.x * q.x)
        yy = (q.y * q.y)
        zz = (q.z * q.z)

        ww - xx - yy + zz
    else
        throw(BoundsError(r,i))
    end
end

# This speeds up some StaticArray methods (such as conversion to RotMatrix)
function Base.convert(::Type{Tuple}, q::Quat)
    ww = (q.w * q.w)
    xx = (q.x * q.x)
    yy = (q.y * q.y)
    zz = (q.z * q.z)
    xy = (q.x * q.y)
    zw = (q.w * q.z)
    xz = (q.x * q.z)
    yw = (q.y * q.w)
    yz = (q.y * q.z)
    xw = (q.w * q.x)

    # initialize rotation part
    return (ww + xx - yy - zz,
            2 * (xy + zw),
            2 * (xz - yw),
            2 * (xy - zw),
            ww - xx + yy - zz,
            2 * (yz + xw),
            2 * (xz + yw),
            2 * (yz - xw),
            ww - xx - yy + zz)
end

function Base.:*(q::Quat, v::StaticVector)
    if length(v) != 3
        throw("Dimension mismatch: cannot rotate a vector of length $(length(v))")
    end

    qo = (-q.x * v[1] - q.y * v[2] - q.z * v[3],
           q.w * v[1] + q.y * v[3] - q.z * v[2],
           q.w * v[2] - q.x * v[3] + q.z * v[1],
           q.w * v[3] + q.x * v[2] - q.y * v[1])

    T = promote_type(eltype(q), eltype(v))

    return similar_type(v, T)(-qo[1] * q.x + qo[2] * q.w - qo[3] * q.z + qo[4] * q.y,
                              -qo[1] * q.y + qo[2] * q.z + qo[3] * q.w - qo[4] * q.x,
                              -qo[1] * q.z - qo[2] * q.y + qo[3] * q.x + qo[4] * q.w)
end

# TODO consider Ac_mul_B, etc...
function Base.:*(q1::Quat, q2::Quat)
    Quat(q1.w*q2.w - q1.x*q2.x - q1.y*q2.y - q1.z*q2.z,
            q1.w*q2.x + q1.x*q2.w + q1.y*q2.z - q1.z*q2.y,
            q1.w*q2.y - q1.x*q2.z + q1.y*q2.w + q1.z*q2.x,
            q1.w*q2.z + q1.x*q2.y - q1.y*q2.x + q1.z*q2.w)
end

function Base.inv(q::Quat)
    Quat(q.w, -q.x, -q.y, -q.z)
end

@inline Base.eye(::Type{Quat}) = Quat(1.0, 0.0, 0.0, 0.0)
@inline Base.eye{T}(::Type{Quat{T}}) = Quat{T}(one(T), zero(T), zero(T), zero(T))

################################################################################
################################################################################
"""
    immutable SPQuat{T} <: Rotation{T}
    SPQuat(x, y, z)

An `SPQuat` is a 3D rotation matrix represented by the "stereographic projection" of a normalized quaternion (shortened to "SPQuat"), which is
a 3-element parametrization of a unit quaternion Q formed by the intersection of a line from [-1,0,0,0] to Q, with a plane containing the origin and with normal direction [1,0,0,0]. This
is a compact representation of rotations where the derivitives of the rotation matrix's elements w.r.t. the SPQuat parameters are rational functions (making them useful for optimization).

See:

    [1] "A Recipe on the Parameterization of Rotation Matrices for Non-Linear Optimization using Quaternions",
    Terzakis, Culverhouse, Bugmann, Sharma, Sutton,
    MIDAS technical report, 2011
    http://www.tech.plymouth.ac.uk/sme/springerusv/2011/publications_files/Terzakis%20et%20al%202012,%20A%20Recipe%20on%20the%20Parameterization%20of%20Rotation%20Matrices...MIDAS.SME.2012.TR.004.pdf

    Note 1: the singularity (origin) has been moved from [0,0,0,-1] in Ref[1] to [-1,0,0,0], so the 0 rotation quaternion [1,0,0,0] maps to [0,0,0] as opposed of to [1,0,0].
    Note 2: the quaternion rotation ambiguity q = -q means there will be a rotation with ||SPQuat|| <= 1 and an equivilent rotation with ||SPQuat|| > 1.  This package
            uses the solution with ||SPQuat|| <= 1
    Note 3: it is safe to assume that the corresponding matrix is orthogonal/unitary for any input x, y, z.

"""
immutable SPQuat{T} <: Rotation{T}
    x::T
    y::T
    z::T

    # TODO should we enforce norm <= 1?
end

# StaticArrays will take over *all* the constructors and put everything in a tuple...
# but this isn't quite what we mean when we have 3 inputs (not 9).
@inline (::Type{SPQuat}){X,Y,Z}(x::X, y::Y, z::Z) = SPQuat{promote_type(promote_type(X, Y), Z)}(x, y, z)

# These 2 functions are enough to satisfy the entire StaticArrays interface:
@inline (::Type{SPQ}){SPQ <: SPQuat}(t::NTuple{9}) = SPQ(Quat(t))
@inline Base.getindex(spq::SPQuat, i::Integer) = Quat(spq)[i]

@inline function Base.convert{Q <: Quat}(::Type{Q}, spq::SPQuat)
    # Both the sign and norm of the Quat is automatically dealt with in its inner constructor
    return Q(1 - (spq.x*spq.x + spq.y*spq.y + spq.z*spq.z), 2*spq.x, 2*spq.y, 2*spq.z)
end

@inline function Base.convert{SPQ <: SPQuat}(::Type{SPQ}, q::Quat)
    alpha2 = (1 - q.w) / (1 + q.w) # <= 1 since q.w >= 0
    spq = SPQ(q.x * (alpha2 + 1)*0.5,  q.y * (alpha2 + 1)*0.5, q.z * (alpha2 + 1)*0.5)
end

@inline Base.convert(::Type{Tuple}, spq::SPQuat) = Tuple(Quat(spq))

@inline Base.:*(spq::SPQuat, x::StaticVector) = Quat(spq) * x

@inline Base.:*(spq::SPQuat, r::Rotation) = Quat(spq) * r
@inline Base.:*(spq::SPQuat, r::RotMatrix) = Quat(spq) * r
@inline Base.:*(r::Rotation, spq::SPQuat) = r * Quat(spq)
@inline Base.:*(r::RotMatrix, spq::SPQuat) = r * Quat(spq)
@inline Base.:*(spq1::SPQuat, spq2::SPQuat) = Quat(spq1) * Quat(spq2)

@inline Base.inv(spq::SPQuat) = SPQuat(-spq.x, -spq.y, -spq.z)

@inline Base.eye(::Type{SPQuat}) = SPQuat(0.0, 0.0, 0.0)
@inline Base.eye{T}(::Type{SPQuat{T}}) = SPQuat{T}(zero(T), zero(T), zero(T))

# rotation properties
@inline rotation_angle(spq::SPQuat) = rotation_angle(Quat(spq))
@inline rotation_axis(spq::SPQuat) = rotation_axis(Quat(spq))


using Knet: @diff, Param, value, grad

sigm(x) = 1.0 / (1.0 + exp(-x))

mutable struct Layer
    wai::Param
    was::Param
    ba::Param
    wri::Param
    wrs::Param
    br::Param
    wsi::Param
    bs::Param
    state
end

Layer(in_size, out_size) =
begin
    sq = sqrt(in_size)
    wai = Param(randn(in_size, out_size)/sq)
    was = Param(randn(out_size, out_size)/sq)
    ba  = Param(zeros(1, out_size))
    wri = Param(randn(in_size, out_size)/sq)
    wrs = Param(randn(out_size, out_size)/sq)
    br  = Param(zeros(1, out_size))
    wsi = Param(randn(in_size, out_size)/sq)
    bs  = Param(zeros(1, out_size))
    state = zeros(1, out_size)
    layer = Layer(wai, was, ba, wri, wrs, br, wsi, bs, state)
layer
end

(layer::Layer)(in) =
begin
    attention = sigm.(in * layer.wai + layer.state * layer.was + layer.ba)
    remember  = sigm.(in * layer.wri + layer.state * layer.wrs + layer.br)
    short_mem = tanh.(in * layer.wsi + attention .* layer.state + layer.bs)
    layer.state = remember .* short_mem + (1 .- remember) .* layer.state
end

make(in_size, layers, out_size) =
begin
    model = []
    hm_layers = length(layers) +1
    for i in 1:hm_layers
        if     i == 1         layer = Layer(in_size, layers[i])
        elseif i == hm_layers layer = Layer(layers[end], out_size)
        else                  layer = Layer(layers[i-1], layers[i])
        end
        push!(model, layer)
    end
model
end

prop(model, in) =
begin
    for layer in model
        in = layer(in)
    end
    for layer in model
        setfield!(layer, :state, zeros(1, size(getfield(layer, :bs))[end]))
    end
in
end

cross_entropy(out, label) =
begin
    - sum(label .* log.(out))
end


#### Operators for a system with a body and stationary points ####

export temperature_bc_constraint_rhs!, heatflux_bc_constraint_rhs!

function temperature_bc_constraint_rhs!(Ts::ScalarData{N},sys::HeatConduction{NX,NY,N},t::Real) where {NX,NY,N}
    surface_temperature!(sys.ΔTs,sys,t)
    Ts .= sys.ΔTs
    return Ts
end

function heatflux_bc_constraint_rhs!(qs::ScalarData{N},sys::HeatConduction{NX,NY,N},t::Real) where {NX,NY,N}
    fill!(qs,0.0)
    # for now, it's adiabatic only
    return qs
end

# Constraint operators, using stored regularization and interpolation operators
# B₁ᵀ = Eᵀ
function temperature_constraint_force!(out::Nodes{Primal,NX,NY},σ::ScalarData{N},sys::HeatConduction{NX,NY,N}) where {NX,NY,N}
    out .= sys.Rc*σ
    return out
end

function temperature_constraint_force!(out::Nodes{Primal,NX,NY},σ,sys::HeatConduction{NX,NY,0}) where {NX,NY}
    return out
end

function heatflux_constraint_force!(out::Nodes{Primal,NX,NY},τ::ScalarData{N},sys::HeatConduction{NX,NY,N}) where {NX,NY,N}
    @unpack params, normals, Rf, Vb, Vf = sys
    @unpack d, α = params
    fact = α/cellsize(sys)
    #fact = isnothing(d) ? fact : fact/d

    fill!(out,0.0)
    product!(Vb,normals,τ)
    Vf .= Rf*Vb
    divergence!(out,Vf)
    out .*= fact
    return out
end

function heatflux_constraint_force!(out::Nodes{Primal,NX,NY},σ,sys::HeatConduction{NX,NY,0}) where {NX,NY}
    return out
end

#=
function _vel_ns_op_constraint_force!(u::Edges{Primal,NX,NY},τ,sys::HeatConduction{NX,NY,0}) where {NX,NY}
    return u
end

function _vel_ns_op_constraint_force!(u::Edges{Primal,NX,NY},τ::VectorData{N},sys::HeatConduction{NX,NY,N}) where {NX,NY,N}
    u .= sys.Rf*τ
    return u
end
=#

# B₂ = E
function temperature_constraint_op!(out::ScalarData{N},T::Nodes{Primal,NX,NY},sys::HeatConduction{NX,NY,N}) where {NX,NY,N}
    out .= sys.Ec*T
    return out
end

function heatflux_constraint_op!(out::ScalarData{N},T::Nodes{Primal,NX,NY},sys::HeatConduction{NX,NY,N}) where {NX,NY,N}
    @unpack params, normals, Ef, gradTb, gradTf = sys
    @unpack k = params
    fill!(gradTf,0.0)

    grad!(gradTf,T)
    gradTb .= Ef*gradTf
    pointwise_dot!(out,gradTb,normals)
    out .*= -k/cellsize(sys)
    return out
end


# examples/cancellation.jl shows effects of these settings
const POWERSERIES_ZMAX = 1.5
const POWERSERIES_NMAX = 11
const POWERSERIES_ZMAX_PRIME = 2.0
const POWERSERIES_NMAX_PRIME = 13

# GSL implementation faster than SpecialFunctions
import GSL
# This survives large arguments better
besselK0(z) = GSL.sf_bessel_K0_scaled(z)*exp(-z)
besselK1(z) = GSL.sf_bessel_K1_scaled(z)*exp(-z)
besselI0(z) = GSL.sf_bessel_I0_scaled(z)*exp(z)
besselI1(z) = GSL.sf_bessel_I1_scaled(z)*exp(z)

# besselK0(z) = besselk(0, z)
# besselK1(z) = besselk(1, z)
# besselI0(z) = besseli(0, z)
# besselI1(z) = besseli(1, z)

# ========== STOKESLET
"Computes S_{ij} f_j"
function stokeslet(r, f, alpha)
    rsq = r[1]*r[1] + r[2]*r[2]
    zsq = rsq*alpha*alpha
    z = sqrt(zsq)
    K0 = besselK0(z)
    K1 = besselK1(z)
    S1 = (zsq*K0 + z*K1 - 1) / (2*pi*zsq)
    S2 = -(zsq*K0 + 2*z*K1 - 2) / (2*pi*zsq*zsq)
    rdotf = (r[1]*f[1] + r[2]*f[2])
    return f.*S1 .+ r.*(alpha*alpha*S2*rdotf)
end

function stokeslet_kernel_noarrays(r1, r2, alpha)
    rsq = r1*r1 + r2*r2
    zsq = rsq*alpha*alpha
    z = sqrt(zsq)
    K0 = besselK0(z)
    K1 = besselK1(z)
    S1 = (zsq*K0 + z*K1 - 1) / (2*pi*zsq)
    S2 = -(zsq*K0 + 2*z*K1 - 2) / (2*pi*zsq*zsq)
    A11 = S1 + alpha^2*S2*r1*r1
    A12 = alpha^2*S2*r1*r2
    A21 = A12
    A22 = S1 + alpha^2*S2*r2*r2
    return A11, A12, A21, A22
end

# ========== PRESSURE(LET)
"Computes p_i f_i"
function pressure(r, f, alpha)
    rsq = r[1]*r[1] + r[2]*r[2]
    rdotf = (r[1]*f[1] + r[2]*f[2])
    return rdotf/(2*pi*rsq)
end

# ========= STRESSLET
"Computes T_{ijk} f_i n_k"
function stresslet(r::Array{Float64}, f::Array{Float64},
                   n::Array{Float64}, alpha::Float64)
    rsq = r[1]*r[1] + r[2]*r[2]
    rnorm = sqrt(rsq)
    z = rnorm*alpha
    (T1, T2, T3) = Ti(z)

    rdotf = r[1]*f[1] + r[2]*f[2]
    rdotn = r[1]*n[1] + r[2]*n[2]
    fdotn = f[1]*n[1] + f[2]*n[2]

    alpha2 = alpha*alpha
    alpha4 = alpha2*alpha2

    u1 = alpha2*T1*(r*fdotn + f*rdotn + n*rdotf)
    u2 = alpha4*T2*rdotf*rdotn*r
    u3 = alpha2*T3*r*fdotn;

    return u1 + u2 + u3
end

"Computes K_{ji} = T_{ijk} n_k"
function dblkernel(r::Array{Float64},
                   n::Array{Float64},
                   alpha::Float64 )
    K = Array{Float64}(undef, 2,2)
    dblkernel!(r, n, alpha, K)
    return K
end

function dblkernel!(r::Array{Float64},
                    n::Array{Float64},
                    alpha::Float64,
                    K::Array{Float64})
    K11, K12, K21, K22 = dblkernel_noarrays(r[1], r[2], n[1], n[2], alpha)
    K[1,1] = K11
    K[1,2] = K12
    K[2,1] = K21
    K[2,2] = K22
end

# Compute kernel without passing around arrays
function dblkernel_noarrays(r1, r2, n1, n2, alpha)
    rsq = r1*r1 + r2*r2
    rnorm = sqrt(rsq)
    z = rnorm*alpha
    (T1, T2, T3) = Ti(z)
    alpha2 = alpha*alpha
    alpha4 = alpha2*alpha2
    T1 *= alpha2
    T2 *= alpha4
    T3 *= alpha2
    rdotn = r1*n1 + r2*n2
    K11 = T1*( n1*r1 + n1*r1 + rdotn ) +
        T2*r1*r1*rdotn +
        T3*n1*r1
    K21 = T1*( n2*r1 + n1*r2  ) +
        T2*r1*r2*rdotn +
        T3*n1*r2
    K12 = T1*( n1*r2 + n2*r1  ) +
        T2*r2*r1*rdotn +
        T3*n2*r1
    K22 = T1*( n2*r2 + n2*r2 + rdotn ) +
        T2*r2*r2*rdotn +
        T3*n2*r2
    K11, K12, K21, K22
end

"""
Computes K_^L{ji}=T_^L{ijk} n_k
# Returns: KL
"""
function dblkernel_log(r::Array{Float64},
                       n::Array{Float64},
                       alpha::Float64 )
    KL = Array{Float64}(undef, 2,2)
    dblkernel_log!(r, n, alpha, KL)
    return KL
end

function dblkernel_log!(r::Array{Float64},
                        n::Array{Float64},
                        alpha::Float64,
                        KL::Array{Float64,2} )
    rsq = r[1]*r[1] + r[2]*r[2]
    rnorm = sqrt(rsq)
    z = rnorm*alpha
    (T1S, T1L, T2S, T2L, T3S, T3L) = Ti_split(z)
    alpha2 = alpha*alpha
    alpha4 = alpha2*alpha2
    T1L *= alpha2
    T2L *= alpha4
    T3L *= alpha2
    rdotn = r[1]*n[1] + r[2]*n[2]
    for i=1:2
        for j=1:2
            KL[j + 2*(i-1)] =
                T1L*( n[j]*r[i] + n[i]*r[j] + (i==j)*rdotn ) +
                T2L*r[i]*r[j]*rdotn +
                T3L*n[i]*r[j]
        end
    end
end

"""
Computes T_{ijk} f_i n_k, but with kernel-splitting
# Returns: (uS, uL, uC, uQ)
such that
u = uS + uL*log(r) + uC*r_k*n_k/r^2 + uQ*f_i*r_i*r_j*r_k*n_k/r^4
"""
function stresslet_split(r, f, n, alpha)
    rnorm = sqrt(r[1]*r[1] + r[2]*r[2])
    z = rnorm*alpha
    (T1S, T1L, T2S, T2L, T3S, T3L) = Ti_split(z)
    rdotf = r[1]*f[1] + r[2]*f[2]
    rdotn = r[1]*n[1] + r[2]*n[2]
    fdotn = f[1]*n[1] + f[2]*n[2]

    vec1 = r*fdotn + f*rdotn + n*rdotf
    vec2 = r*(rdotf*rdotn)
    vec3 = r*fdotn

    alpha2 = alpha*alpha
    alpha4 = alpha2*alpha2

    uS = alpha2*(T1S*vec1 + alpha2*T2S*vec2 + T3S*vec3)
    uL = alpha2*(T1L*vec1 + alpha2*T2L*vec2 + T3L*vec3)
    uS = uS + log(alpha)*uL # Split on log(r)
    uC = alpha2/(8*pi)*r*rdotf
    uQ = -1/pi

    return (uS, uL, uC, uQ)
end

function stresslet_split_noarrays(r1, r2, f1, f2, n1, n2, alpha)
    rnorm = sqrt(r1*r1 + r2*r2)
    z = rnorm*alpha
    (T1S, T1L, T2S, T2L, T3S, T3L) = Ti_split(z)
    rdotf = r1*f1 + r2*f2
    rdotn = r1*n1 + r2*n2
    fdotn = f1*n1 + f2*n2

    vecA1 = r1*fdotn + f1*rdotn + n1*rdotf
    vecA2 = r2*fdotn + f2*rdotn + n2*rdotf
    vecB1 = r1*(rdotf*rdotn)
    vecB2 = r2*(rdotf*rdotn)
    vecC1 = r1*fdotn
    vecC2 = r2*fdotn

    alpha2 = alpha*alpha
    alpha4 = alpha2*alpha2

    uS1 = alpha2*(T1S*vecA1 + alpha2*T2S*vecB1 + T3S*vecC1)
    uS2 = alpha2*(T1S*vecA2 + alpha2*T2S*vecB2 + T3S*vecC2)
    uL1 = alpha2*(T1L*vecA1 + alpha2*T2L*vecB1 + T3L*vecC1)
    uL2 = alpha2*(T1L*vecA2 + alpha2*T2L*vecB2 + T3L*vecC2)
    uS1 += log(alpha)*uL1 # Split on log(r)
    uS2 += log(alpha)*uL2 # Split on log(r)
    uC1 = alpha2/(8*pi)*r1*rdotf
    uC2 = alpha2/(8*pi)*r2*rdotf
    uQ = -1/pi

    return (uS1, uS2, uL1, uL2, uC1, uC2, uQ)
end

"""
Stresslet functions
# Returns: (T1, T2, T3)
"""
function Ti(z::Float64)
    if z < POWERSERIES_ZMAX
        (T1S, T1L, T2S, T2L, T3S, T3L) = Ti_split_pow(z)
        T1 = T1S + T1L*log(z)
        T2 = T2S + T2L*log(z) + 1/(8*pi*z^2) - 1/(pi*z^4)
        T3 = T3S + T3L*log(z)
        return (T1, T2, T3)
    else
        return Ti_direct(z)
    end
end

function Ti_direct(z::Float64)
    zsq = z*z
    z4 = zsq*zsq
    z6 = z4*zsq
    K0 = besselK0(z)
    K1 = besselK1(z)
    T1 = -(2*zsq*K0 + (zsq+4)*z*K1 - 4) / (2*pi*z4)
    T2 = (4*zsq*K0 + (zsq+8)*z*K1 - 8) / (pi*z6)
    T3 = (z*K1 - 1) / (2*pi*zsq)
    return (T1, T2, T3)
end

function Ti_prime(z::Float64)
    if z < POWERSERIES_ZMAX_PRIME
        (T1S, T1L, T2S, T2L, T3S, T3L) = Ti_split_pow(z, false)
        (DT1S, DT1L, DT2S, DT2L, DT3S, DT3L) = Ti_split_pow(z, true)
        T1p = DT1S + DT1L*log(z) + T1L / z
        T2p = DT2S + DT2L*log(z) + T2L/z - 2/(8*pi*z^3) + 4/(pi*z^5)
        T3p = DT3S + DT3L*log(z) + T3L/z
        return (T1p, T2p, T3p)
    else
        return Ti_direct_prime(z)
    end
end

function Ti_direct_prime(z::Float64)
    zsq = z*z
    z4 = zsq*zsq
    z6 = z4*zsq
    K0 = besselK0(z)
    K1 = besselK1(z)
    T1p = ((2*z^2 + 8)*K1)/(z^4*pi) - 8/(z^5*pi) + ((z^2 + 8)*K0)/(2*z^3*pi)
    T2p = 48/(z^7*pi) - ((8*z^2 + 48)*K1)/(z^6*pi) - ((z^2 + 24)*K0)/(z^5*pi)
    T3p = 1/(z^3*pi) - K0/(2*z*pi) - K1/(z^2*pi)
    return (T1p, T2p, T3p)
end

"""
Split of stresslet functions
# Returns: (T1S, T1L, T2S, T2L, T3S, T3L)
"""
function Ti_split(z::Float64)
    if z < POWERSERIES_ZMAX
        return Ti_split_pow(z)
    else
        return Ti_split_direct(z)
    end
end

function Ti_split_prime(z::Float64)
    if z < POWERSERIES_ZMAX_PRIME
        return Ti_split_pow(z,true)
    else
        return Ti_split_direct(z,true)
    end
end

function Ti_split_direct(z::Float64,ifprime::Bool=false)

    if !ifprime

        K0 = besselK0(z)
        K1 = besselK1(z)
        I0 = besselI0(z)
        I1 = besselI1(z)

        K0S = K0 + I0*log(z)
        K1S = K1 - I1*log(z) - 1/z

        T1S = -(2*z*K0S + (z^2+4)*K1S + z) / (2*pi*z^3)
        T1L = (2*z*I0 - (z^2+4)*I1) / (2*pi*z^3)

        T2S = (32*z*K0S + 8*(z^2+8)*K1S  - z*(z^2-16)) / (8*pi*z^5)
        T2L = ((z^2+8)*I1 - 4*z*I0) / (pi*z^5)

        T3S = K1S/(2*pi*z)
        T3L = I1/(2*pi*z)
        return (T1S, T1L, T2S, T2L, T3S, T3L)

    else

        logz = log(z)
        z2 = z*z
        z3 = z*z2
        z4 = z*z3
        z5 = z*z4
        z6 = z*z5
        z7 = z*z6

        K0 = besselK0(z)
        K1 = besselK1(z)
        I0 = besselI0(z)
        I1 = besselI1(z)

        DT1S = (-16 +
                z*(z*(8+z2)*K0+4*(4+z2)*K1-(4+z2)*I1*(4*logz-1)+
                   z*I0*((8+z2)*logz-2)))/(2*pi*z5)
        DT2S = (192-16*z2+z4+
                4*z*(-z*(24+z2)*K0-8*(6+z2)*K1+
                     I1*(-8-z2+8*(6+z2)*logz)-
                     z*I0*(-4+(24+z2)*logz)))/(4*pi*z7)
        DT3S = (2-I1*(z-2*z*logz)-
                z2*(2*K1/z+K0+I0*logz))/(2*pi*z3)

        DT1L = (-2*z*(8 + z2)*I0 + 8*(4 + z2)*I1)/(4*pi*z4)
        DT2L = (z*(24 + z2)*I0 - 8*(6 + z2)*I1)/(pi*z6)
        DT3L = (I0-2*I1/z)/(2*pi*z)

        return (DT1S, DT1L, DT2S, DT2L, DT3S, DT3L)

    end

end

"""
Direct power series for split stresslet functions
# Returns: (T1S, T1L, T2S, T2L, T3S, T3L)
"""
function Ti_split_pow(z::Float64,ifprime::Bool=false)
    if !ifprime
        # Outputs
        T1S = 0.0
        T1L = 0.0
        T2S = 0.0
        T2L = 0.0
        T3S = 0.0
        T3L = 0.0
        # Add n=0 term
        log2 = log(2)
        k1L0 = 0.5
        k1S0 = (-2*log2-(-2*eulergamma+1.0)) * 0.25
        T3S += k1S0
        T3L += k1L0
        # Recursion variables, init at n=1
        n=1
        z2 = z*z
        z2nm2 = 1; # z^(2*n-2)
        z2n = z2; # z^(2*n)
        digamma_np1 = -eulergamma + 1.0 # digamma(n+1)
        digamma_np2 = -eulergamma + 1.5
        ifact_n = 1.0 # 1/factorial(n)
        ifact_np1 = 1.0/2
        i4n = 1.0/4 # 1/4^n
        i4np1 = 1.0/16
        k1Snm1 = k1S0 # k1S(n-1)
        k1Lnm1 = k1L0 # k1L(n-1)
        ## First iteration unrolled
        # k0, k1
        k0n_denom = i4n*ifact_n*ifact_n
        k1n_denom = i4np1*ifact_n*ifact_np1
        k0Ln = (-1)*k0n_denom
        k0Sn = (digamma_np1 + log2)*k0n_denom
        k1Ln = (2)*k1n_denom
        k1Sn = (-2*log2-digamma_np1-digamma_np2)*k1n_denom
        # T1
        t1Sn = 2*k0Sn + 4*k1Sn + k1Snm1
        t1Ln = 2*k0Ln + 4*k1Ln + k1Lnm1
        T1S += t1Sn*z2nm2
        T1L += t1Ln*z2nm2
        # T3
        T3S += k1Sn*z2n
        T3L += k1Ln*z2n
        # Update recursions
        n += 1
        inp1 = 1/(n+1)
        z2n *= z2
        z2nm2 = z2
        z2nm4 = 1
        (digamma_np1, digamma_np2) = (digamma_np2, digamma_np2+inp1)
        (ifact_n, ifact_np1) = (ifact_np1, ifact_np1*inp1)
        (i4n, i4np1) = (i4np1, i4np1*0.25)
        k1Snm1 = k1Sn
        k1Lnm1 = k1Ln
        ## Loop
        while n<POWERSERIES_NMAX
            # k0, k1
            k0n_denom = i4n*ifact_n*ifact_n
            k1n_denom = i4np1*ifact_n*ifact_np1
            k0Ln = (-1)*k0n_denom
            k0Sn = (digamma_np1 + log2)*k0n_denom
            k1Ln = (2)*k1n_denom
            k1Sn = (-2*log2-digamma_np1-digamma_np2)*k1n_denom
            # T1
            t1Sn = 2*k0Sn + 4*k1Sn + k1Snm1
            t1Ln = 2*k0Ln + 4*k1Ln + k1Lnm1
            T1S += t1Sn*z2nm2
            T1L += t1Ln*z2nm2
            # T2
            t2Sn = 4*k0Sn + 8*k1Sn + k1Snm1
            t2Ln = 4*k0Ln + 8*k1Ln + k1Lnm1
            T2S += t2Sn*z2nm4
            T2L += t2Ln*z2nm4
            # T3
            T3S += k1Sn*z2n
            T3L += k1Ln*z2n
            # Update recursions
            n += 1
            inp1 = 1/(n+1)
            (z2nm4, z2nm2, z2n) = (z2nm2, z2n, z2n*z2)
            (digamma_np1, digamma_np2) = (digamma_np2, digamma_np2+inp1)
            (ifact_n, ifact_np1) = (ifact_np1, ifact_np1*inp1)
            (i4n, i4np1) = (i4np1, i4np1*0.25)
            k1Snm1 = k1Sn
            k1Lnm1 = k1Ln
        end
        # Multiply with constant factors
        T1S *= -1/(2*pi)
        T1L *= -1/(2*pi)
        T2S *= 1/pi
        T2L *= 1/pi
        T3S *= 1/(2*pi)
        T3L *= 1/(2*pi)
        return (T1S, T1L, T2S, T2L, T3S, T3L)
    else
        # Outputs
        T1S = 0.0
        T1L = 0.0
        T2S = 0.0
        T2L = 0.0
        T3S = 0.0
        T3L = 0.0
        # Add n=0 term (derivative is zero)
        log2 = log(2)
        k1L0 = 0.5
        k1S0 = (-2*log2-(-2*eulergamma+1.0)) * 0.25
        #T3S += k1S0
        #T3L += k1L0
        # Recursion variables, init at n=1
        n=1
        z2 = z*z
        z2nm3 = 0; # z^(2*n-3)*(2n-2)
        z2nm1 = z*2; # z^(2*n-1)*2n
        digamma_np1 = -eulergamma + 1.0 # digamma(n+1)
        digamma_np2 = -eulergamma + 1.5
        ifact_n = 1.0 # 1/factorial(n)
        ifact_np1 = 1.0/2
        i4n = 1.0/4 # 1/4^n
        i4np1 = 1.0/16
        k1Snm1 = k1S0 # k1S(n-1)
        k1Lnm1 = k1L0 # k1L(n-1)
        ## First iteration unrolled
        # k0, k1
        k0n_denom = i4n*ifact_n*ifact_n
        k1n_denom = i4np1*ifact_n*ifact_np1
        k0Ln = (-1)*k0n_denom
        k0Sn = (digamma_np1 + log2)*k0n_denom
        k1Ln = (2)*k1n_denom
        k1Sn = (-2*log2-digamma_np1-digamma_np2)*k1n_denom
        # T1 (constant ... derivative is zero)
        #t1Sn = 2*k0Sn + 4*k1Sn + k1Snm1
        #t1Ln = 2*k0Ln + 4*k1Ln + k1Lnm1
        #T1S += t1Sn*z2nm3
        #T1L += t1Ln*z2nm3
        # T3
        T3S += k1Sn*z2nm1
        T3L += k1Ln*z2nm1
        # Update recursions
        n += 1
        inp1 = 1/(n+1)
        z2nm1 *= z2*(n/(n-1.0))
        z2nm3 = z*2
        z2nm5 = 0
        (digamma_np1, digamma_np2) = (digamma_np2, digamma_np2+inp1)
        (ifact_n, ifact_np1) = (ifact_np1, ifact_np1*inp1)
        (i4n, i4np1) = (i4np1, i4np1*0.25)
        k1Snm1 = k1Sn
        k1Lnm1 = k1Ln
        ## Loop
        while n<POWERSERIES_NMAX_PRIME
            # k0, k1
            k0n_denom = i4n*ifact_n*ifact_n
            k1n_denom = i4np1*ifact_n*ifact_np1
            k0Ln = (-1)*k0n_denom
            k0Sn = (digamma_np1 + log2)*k0n_denom
            k1Ln = (2)*k1n_denom
            k1Sn = (-2*log2-digamma_np1-digamma_np2)*k1n_denom
            # T1
            t1Sn = 2*k0Sn + 4*k1Sn + k1Snm1
            t1Ln = 2*k0Ln + 4*k1Ln + k1Lnm1
            T1S += t1Sn*z2nm3
            T1L += t1Ln*z2nm3
            # T2
            t2Sn = 4*k0Sn + 8*k1Sn + k1Snm1
            t2Ln = 4*k0Ln + 8*k1Ln + k1Lnm1
            T2S += t2Sn*z2nm5
            T2L += t2Ln*z2nm5
            # T3
            T3S += k1Sn*z2nm1
            T3L += k1Ln*z2nm1
            # Update recursions
            n += 1
            inp1 = 1/(n+1)
            (z2nm5, z2nm3, z2nm1) = (z2nm3, z2nm1, z2nm1*z2*(n/(n-1.0)))
            (digamma_np1, digamma_np2) = (digamma_np2, digamma_np2+inp1)
            (ifact_n, ifact_np1) = (ifact_np1, ifact_np1*inp1)
            (i4n, i4np1) = (i4np1, i4np1*0.25)
            k1Snm1 = k1Sn
            k1Lnm1 = k1Ln
        end
# Multiply with constant factors
        T1S *= -1/(2*pi)
        T1L *= -1/(2*pi)
        T2S *= 1/pi
        T2L *= 1/pi
        T3S *= 1/(2*pi)
        T3L *= 1/(2*pi)
        return (T1S, T1L, T2S, T2L, T3S, T3L)
    end
end


"""
    psf(Environment[,cent])

Calculate frequency-domain point spread function using the Environment struct to access
steeringvectors. Optionally, supply the index where the psf is centered, default is (N/2)+1.
"""
function psf(E::Environment,cent::Int64=floor(Int,E.N/2)+1)
    @unpack steeringvec,M,N,Nf,fn = E
    p = Array{Float64, 2}(undef, N, Nf)
    @views @inbounds for j in 1:Nf
        AeroAcoustics.psf_col!(p[:,j],steeringvec.arr[:,:,j],cent)
    end
    return FreqArray(p,fn)
end

"""
    psf_col!(p,steeringvec,cent)

Calculate single frequency point spread function as a column vector.
"""
function psf_col!(p::A,steeringvec::B,cent::Int64) where {T<:AbstractFloat, N, A<:AbstractArray{T,1}, B<:AbstractArray{Complex{T},N}}
    M = size(steeringvec,1)
    @views p .= M^2 .*abs2.(steeringvec'*steeringvec[:,cent])
end


"""
A struct that represents an agent storing it's accuracy and balance.
The accurancy is set up once during the agent creation and remain the same during the simulation.
The balance is updated accoring to agent's perfromance in the curation process.
"""
mutable struct Agent{TI<:Integer, TF<:Float64}
    accuracy::TI
    balance::TF
end

# constructor
Agent(a::Integer) = Agent(min(100,a), 50.0)

"""
Creates a group of X (count) agents all of them having the same accuracy (acc)
"""
function setupAgentsWithFixedAccuracy(count, acc)
    Agent[Agent(acc) for i in 1:count]
end

"""
Creates a group of X (count) agents all with accuracy distributed according to the
normal distribution with passed parameters (mean and std)
"""
function setupRandomAgents(count, mean, std)
    agents = Agent[Agent(trunc(Int, rand(TruncatedNormal(mean, std, 0, 100)))) for i in 1:count]
end

function setupMixedAgents(lowCount, lowAcc, highCount, highAcc)
    agents = vcat(setupAgentsWithFixedAccuracy(lowCount, lowAcc), setupAgentsWithFixedAccuracy(highCount, highAcc))
end

"""
This function represents the perception of an item by the agent.
Agents with lower accuracy tend to overrate an item.
This variation of the function rates the item with a fixed positive bias.
"""
function evaluateCandidateByAgentFixed(candidate, agent)
    evaluation = candidate + (100-agent.accuracy)/3
end

"""
This function represents the perception of an item by the agent.
Agents with lower accuracy tend to overrate an item.
This variation of the function rates the item with a randomized bias based on the normal distribution.
"""
function evaluateCandidateByAgent(candidate, agent)
    if (agent.accuracy == 100)
        candidate
    else
        evaluation = rand(Normal(candidate, (100-agent.accuracy)/3))
        candidate + abs(evaluation - candidate)
    end
end


module RVFLMod

using DataFrames
using LinearAlgebra
using Random
using ..DataMod

export RVFL, enhanced, radbas

randomrange(rng::AbstractRNG, C::Number, dims::Integer...)::VecOrMat{Float64} =
    2C * rand(rng, dims...) .- C
enames(n::Integer)::Vector{Symbol} = Symbol.('e', 1:n)

# RVFL Input layer
struct RVFL
    m::Integer
    A::Matrix{Float64}
    b::Vector{Float64}
    g::Function

    RVFL(in::Integer, enh::Integer, rng::AbstractRNG = MersenneTwister(); range::Float64 = 1.0, act::Function=radbas) =
        new(enh, randomrange(rng, range, enh, in), randomrange(rng, range, enh), act)
end

enhanced(m::RVFL, x::Union{Vector{Float64}, SubArray{Float64}}) =
    [x; map(m.g, [dot(m.A[i, :], x) for i in 1:m.m] .+ m.b)]

enhanced(m::RVFL, x::Matrix{Float64}) =
    mapslices(row -> enhanced(m, row), Matrix(x), dims = 2)

enhanced(m::RVFL, xdf::DataFrame) =
    DataFrame(enhanced(m, Matrix(xdf)), [propertynames(xdf); enames(m.m)])

enhanced(m::RVFL, d::Data) =
    Data(d, hcat(enhanced(m, d.xdf), d.ydf), [propertynames(d.xdf); enames(m.m)])

# Activation functions
radbas(x::Float64) = exp(-x^2)

end

using .RVFLMod


using ControlSystems
using Statistics

"""
    sv, sx, sigma_se, max_sx, max_sy = simulateMS(v, M=16, mtf=zpk(1,0,1,1), d=0, dw=[1..], sx0=[0..])

Simulate the vector-based Mismatch-Shaping for a multi-element DAC.
Inputs:
- `v`:        A vector of the digital input values. `v` is in (-`M`:2:`M`) if `dw`=[1..]
              otherwise `v` is in [-sum(`dw`),sum(`dw`)].
- `M`:        The number of elements.
- `mtf`:	  The mismatch-shaping transfer function, given in zero-pole form.
- `d`:        Dither uniformly distributed in [-`d`,`d`] is added to sy.
- `dw`:       A vector of dac element weights
- `sx0`:      A matrix whose columns are the initial states of the ESL.
Outputs:
- `sv`:       An MxN matrix whose columns are the selection vectors.
- `sx`:       An orderxM matrix containing the final state of the ESL.
- `sigma_se`: The rms value of the selection error.
- `max_sx`:   The maximum absolute value of the state for all modulators.
- `max_sy`:   The maximum absolute value of the input to the VQ.

For a brief description of the theory of mismatch-shaping DACs, see
R. Schreier and B. Zhang "Noise-shaped multibit D/A convertor employing
unit elements," Electronics Letters, vol. 31, no. 20, pp. 1712-1713,
Sept. 28 1995.
"""
function simulateMS(v, M=16, mtf=zpk([1], [0], 1, 1), d=0, dw=[], sx0=[])
    order = length(mtf.matrix[1].p)
    if isempty(sx0)
        sx0 = zeros(order, M)
    end
    if isempty(dw)
        dw = ones(M)
    end

    # B/A = MTF-1
    mtf = tf(mtf)
    num = mtf.matrix[1].num[end:-1:0]
    den = mtf.matrix[1].den[end:-1:0]
    A = real(-den[2:order+1])
    B = real(num[2:order+1] + A)

    N = length(v)
    sv = zeros(Int, M, N)

    sx = sx0'
    max_sx = maximum(abs.(sx))
    max_sy = 0
    sum_se2 = 0

    for i = 1:N
        # Compute the sy vector.
        sy = sx * B
        # Normalize sy for a minimum value of zero.
        sy = sy .- minimum(sy)
        # Pick the elements that have the largest desired_usage (sy) values.
        sv[:,i] = selectElement(v[i], sy + d * (2*rand(M) .- 1), dw)
        # Compute the selection error.
        se = sv[:,i] - sy
        # Compute the new sx matrix
        sxn = sx * A + se
        sx = [sxn sx[:,1:order-1]]
        # Keep track of some statistics.
        sum_se2 = sum_se2 + sum((se .- mean(se)).^2)
        max_sx = max(max_sx, maximum(abs.(sxn)))
        max_sy = max(max_sy, maximum(abs.(sy)))
    end
    sigma_se = sqrt(sum_se2 / (M*N))

    return sv, sx', sigma_se, max_sx, max_sy
end


using ModelingToolkit, DiffEqBase, LinearAlgebra

# Define some variables
@parameters t x y
@variables u(..)

ConstrainedEquation([x ~ 0,y < 1/2], u(t,x,y) ~ x + y^2)


#push!(LOAD_PATH, "/home/caseykneale/Desktop/ChemometricsTools/ChemometricsTools.jl/");
using ChemometricsTools
#View the data in the package space
ChemometricsToolsDatasets()
#Load the data from package-space
Iris = convert(Matrix, ChemometricsToolsDataset("iris.csv") ) ;
#Seperate data from labels
X = convert( Array{Float64,2}, Iris[:, 1 : ( end - 1 ) ] );
Y = Iris[:, end ];
#One hot encode the class vector
Encoding = LabelEncoding(Y);
YHot = ColdToHot(Y, Encoding);
#Shuffle the set, because the classes are ordered
Shuffle!( X, YHot );
#Split the data into a training and hold out set
( (TrainX, TrainY), (TestX, TestY) ) = SplitByProportion( X, YHot, 0.7);

#Let the shoot out begin!
@time knn = KNN(TrainX, TrainY, "euclidean")
@time lda = LDA(TrainX, TrainY; Factors = 2)
@time ldagd = GaussianDiscriminant( lda, TrainX, TrainY; Factors = 2 )
@time mnsr = MultinomialSoftmaxRegression(TrainX, TrainY; maxiters = 2000, LearnRate = 1e-3, L2 = 1e4)
@time gnb = GaussianNaiveBayes(TrainX, TrainY)
@time ct = ClassificationTree(TrainX, TrainY;gainfn = entropy, maxdepth = 8, minbranchsize = 5)
@time rf = RandomForest(TrainX, TrainY, :classification; gainfn = entropy, trees = 100,
                        maxdepth = 10,  minbranchsize = 5,
                        samples = 0.7, maxvars = nothing)
@time simca = SIMCA(TrainX, TrainY; VarianceExplained = 0.95, Quantile = 0.999)

#Let's evaluate them on the training set then the hold out - see how they do!
models = ( knn, ldagd, mnsr, gnb, ct, rf, simca );
modelnames = ( :knn, :ldagd, :mnsr, :gnb, :ct, :rf, :SIMCA );
TrainingF1s = Dict()
for ( name, model ) in zip(modelnames, models)
    Predicted = model(TrainX)
    TrainingF1s[name] = MulticlassStats( Predicted, TrainY, Encoding )[1][:Global]["FMeasure"]
end

TrainingF1s

HoldOutF1s = Dict()
for ( name, model ) in zip(modelnames, models)
    HoldOutF1s[name] = MulticlassStats( model(TestX), TestY, Encoding )[1][:Global]["FMeasure"]
end

HoldOutF1s

#Let's check specifically the random forest classification model for the virginica class
RFStats = MulticlassStats( rf(TestX), TestY, Encoding );
RFStats[2]["virginica"]

StatsToCSVs(RFStats, "/home/caseykneale/Desktop/", "RFStats")
StatsToLaTeX(RFStats, "/home/caseykneale/Desktop/", "RFStats"; maxcolumns = 3)


abstract type DiffPrecSupport end
struct ANTSupport <: DiffPrecSupport end
struct BootStdSupport <: DiffPrecSupport end
struct BootMaxSupport <: DiffPrecSupport end   # not implemented yet
struct DTraceValidationSupport <: DiffPrecSupport end

####################################
#
#   support estimator
#
####################################

function __initSupport(
    Sx::Symmetric, Sy::Symmetric, X, Y)

    p = size(X)
    np = div((p + 1)*p, 2)

    estimSupp = Vector{Vector{Int64}}(undef, np)

    # first stage
    λ = 1.01 * quantile(Normal(), 1. - 0.1 / (p * (p+1)))
    H = Symmetric( skron(Sx, Sy) )
    b = svec(Sx) - svec(Sy)
    x1 = reducedDiffEstimation(H, b, Sx, Sy, X, Y, λ)
    S1 = getReducedSupport(x1)

    # second stage
    ind = 0
    for ci=1:p
      for ri=ci:p
        indElem = sub2indLowerTriangular(p, ri, ci)
        # x2 = invQSymHessian(Sx, Sy, j, X, Y, λ)
        # S2 = getSupport(x2, p)
        x2 = invHessianReduced(H, indElem, Sx, Sy, X, Y, λ)
        S2 = getReducedSupport(x2)
        ind = ind + 1
        estimSupp[ind] = (1:np)[S1 .| S2]
      end
    end

    estimSupp
end

function __threshold(Δpoint::Vector{DiffPrecResultNormal}, p::Int, τ0::Float64)
    Δ = zeros(Float64, p, p)
    it = 0
    for col=1:p
        for row=col:p
            it = it + 1
            τ = τ0 * Δpoint[it].std
            v = abs(Δpoint[it].p) > τ ? Δpoint[it].p : 0.
            if row == col
                Δ[row, row] = v
            else
                Δ[row, col] = v
                Δ[col, row] = v
            end
        end
    end
    sparse(Δ)
end

function __threshold(Δpoint::Vector{DiffPrecResultBoot}, p::Int, τ0::Float64)
    Δ = zeros(Float64, p, p)
    it = 0
    for col=1:p
        for row=col:p
            it = it + 1
            τ = τ0 * std( Δpoint[it].boot_p )
            v = abs(Δpoint[it].p) > τ ? Δpoint[it].p : 0.
            if row == col
                Δ[row, row] = v
            else
                Δ[row, col] = v
                Δ[col, row] = v
            end
        end
    end
    sparse(Δ)
end

function __getΔ(::ANTSupport, X, Y; estimSupport::Union{Array{BitArray},Nothing}=nothing)
    nx, p = size(X)
    ny    = size(Y, 1)

    Sx = Symmetric( cov(X) )
    Sy = Symmetric( cov(Y) )

    eS = estimSupport === nothing ? __initSupport(Sx, Sy, X, Y) : estimSupport
    Δpoint = Array{DiffPrecResultNormal}(undef, div((p+1)*p, 2))

    it = 0
    for col=1:p
        for row=col:p
            it = it + 1
            indS = getLinearSupport(row, col, eS[it])

            Δpoint[it] = estimate(SymmetricOracleNormal(), Sx, Sy, X, Y, row, col, indS)
        end
    end

    Δpoint, eS
end

function supportEstimate(::ANTSupport, X, Y, τ0::Float64; estimSupport::Union{Array{BitArray},Nothing}=nothing)
    p = size(X, 2)

    Δpoint, eS = __getΔ(ANTSupport(), X, Y; estimSupport=estimSupport)
    Δsupp = __threshold(Δpoint, p, τ0)

    Δsupp, Δpoint, eS
end

function supportEstimate(::ANTSupport, X, Y, τArr::Vector{Float64}; estimSupport::Union{Array{BitArray},Nothing}=nothing)
    p = size(X, 2)

    Δpoint, eS = __getΔ(ANTSupport(), X, Y; estimSupport=estimSupport)

    Δsupp = Array{SparseMatrixCSC{Float64,Int64}}(undef, length(τArr))
    for i=1:length(τArr)
        Δsupp[i] = __threshold(Δpoint, p, τArr[i])
    end

    Δsupp, Δpoint, eS
end

function __getΔ(::BootStdSupport, X, Y; estimSupport::Union{Array{BitArray},Nothing}=nothing)
    nx, p = size(X)
    ny    = size(Y, 1)

    Sx = Symmetric( cov(X) )
    Sy = Symmetric( cov(Y) )

    eS = estimSupport === nothing ? __initSupport(Sx, Sy, X, Y) : estimSupport
    Δpoint = Array{DiffPrecResultBoot}(undef, div((p+1)*p, 2))

    it = 0
    for col=1:p
        for row=col:p
            it = it + 1
            indS = getLinearSupport(row, col, eS[it])

            Δpoint[it] = estimate(SymmetricOracleBoot(), Sx, Sy, X, Y, row, col, indS)
        end
    end

    Δpoint, eS
end

function supportEstimate(::BootStdSupport, X, Y, τ0::Float64; estimSupport::Union{Array{BitArray},Nothing}=nothing)
    p = size(X, 2)

    Δpoint, eS = __getΔ(BootStdSupport(), X, Y; estimSupport=estimSupport)
    Δsupp = __threshold(Δpoint, p, τ0)

    Δsupp, Δpoint, eS
end

function supportEstimate(::BootStdSupport, X, Y, τArr::Vector{Float64}; estimSupport::Union{Array{BitArray},Nothing}=nothing)
    p = size(X, 2)

    Δpoint, eS = __getΔ(BootStdSupport(), X, Y; estimSupport=estimSupport)

    Δsupp = Array{SparseMatrixCSC{Float64,Int64}}(undef, length(τArr))
    for i=1:length(τArr)
        Δsupp[i] = __threshold(Δpoint, p, τArr[i])
    end

    Δsupp, Δpoint, eS
end

####################################
#
#   support estimator  -- DTrace
#
####################################

# Λarr should be in descending order
function supportEstimate(::DTraceValidationSupport,
    Sx::Symmetric, Sy::Symmetric,
    SxValid::Symmetric, SyValid::Symmetric,
    Λarr::Vector{Float64},
    options=CDOptions())

    numΛ = length(Λarr)
    eΔarr = Vector{SparseMatrixCSC{Float64,Int64}}(undef, numΛ)
    loss2arr = Vector{Float64}(undef, numΛ)
    lossInfarr = Vector{Float64}(undef, numΛ)

    f = CDDirectDifferenceLoss(Sx, Sy)
    x = SymmetricSparseIterate(f.p)
    g = ProxL1(Λarr[1])
    coordinateDescent!(x, f, g, options)
    eΔarr[1] = sparse(Matrix(x))
    loss2arr[1] = CovSel.diffLoss(SxValid, x, SyValid, 2)
    lossInfarr[1] = CovSel.diffLoss(SxValid, x, SyValid, Inf)

    opt = CDOptions(;
        maxIter=options.maxIter,
        optTol=options.optTol,
        randomize=options.randomize,
        warmStart=true,
        numSteps=options.numSteps)

    for i=2:numΛ
        g = ProxL1(Λarr[i])
        coordinateDescent!(x, f, g, opt)
        eΔarr[i] = sparse(Matrix(x))
        loss2arr[i] = CovSel.diffLoss(SxValid, x, SyValid, 2)
        lossInfarr[i] = CovSel.diffLoss(SxValid, x, SyValid, Inf)
    end

    # find min loss
    i2 = argmin(loss2arr)
    iInf = argmin(loss2arr)

    eΔarr, i2, iInf, loss2arr, lossInfarr
end


using SolverPeeker
using Test
using JuMP

"""
Initiate some models
"""

# Model with linear constraints and objective
model1 = Model()
@variable(model1, x >= 0, Int)
@variable(model1, y <= 5)
@variable(model1, z)
@objective(model1, Max, 3x+5y-2z)
@constraint(model1, x+2y <= 45)
@constraint(model1, z - y >= 4)
@constraint(model1, 0<= 2z <= 7)

# Model with nothging specified. Should have objective 0.
model2 = Model()

# Objective registered as 0 in model(?)
model_nl = Model()
@variable(model_nl, x, start = 0.0)
@variable(model_nl, y, start = 0.0)
@NLobjective(model_nl, Min, (1 - x)^2 + 100 * (y - x^2)^2)
@constraint(model_nl, x + y == 10)
#@NLobjective(model_nl, Min, sin(x))

@testset "recognize" begin

    """
    Constant objective test
    """
    """
    @test recognize(model1, ObjectiveIsConstant()) == false

    @test recognize(model2, ObjectiveIsConstant()) == true

    # Add variable and constant objective. Should still be true:
    @variable(model2, x)
    @objective(model2, Min, 5)
    @test recognize(model2, ObjectiveIsConstant()) == true
    """

end


module PointMassTest

using Test
using ReactiveMP
using Distributions
using Random
using SpecialFunctions

import ReactiveMP: ∞

@testset "PointMass" begin

    @testset "Real-based PointMass" begin

        scalar = rand()
        dist   = PointMass(scalar)

        @test variate_form(dist) === Univariate
        @test_throws BoundsError dist[2]
        @test_throws BoundsError dist[2, 2]

        @test insupport(dist, scalar)
        @test !insupport(dist, scalar + 1e-6)
        @test !insupport(dist, scalar - 1e-6)

        @test pdf(dist, scalar)        === 1.0
        @test pdf(dist, scalar + 1e-6) === 0.0
        @test pdf(dist, scalar - 1e-6) === 0.0

        @test logpdf(dist, scalar)        === 0.0
        @test logpdf(dist, scalar + 1e-6) === -Inf
        @test logpdf(dist, scalar - 1e-6) === -Inf

        @test_throws MethodError insupport(dist, [ 1.0, 1.0 ])
        @test_throws MethodError insupport(dist, [ 1.0 1.0; 1.0 1.0 ])
        @test_throws MethodError pdf(dist, [ 1.0, 1.0 ])
        @test_throws MethodError pdf(dist, [ 1.0 1.0; 1.0 1.0 ])
        @test_throws MethodError logpdf(dist, [ 1.0, 1.0 ])
        @test_throws MethodError logpdf(dist, [ 1.0 1.0; 1.0 1.0 ])

        @test entropy(dist) === -∞

        @test mean(dist)       === scalar
        @test var(dist)        === 0.0
        @test std(dist)        === 0.0
        @test cov(dist)        === 0.0
        @test precision(dist)  === Inf
        @test ndims(dist)      === 1

        @test_throws ErrorException probvec(dist)
        @test logmean(dist)         === log(scalar)
        @test inversemean(dist)     === inv(scalar)
        @test mirroredlogmean(dist) === log(1.0 - scalar)
        @test loggammamean(dist)    === loggamma(scalar)
    end

    @testset "Vector-based PointMass" begin

        vector = rand(5)
        dist   = PointMass(vector)

        @test variate_form(dist) === Multivariate
        @test dist[2] === vector[2]
        @test dist[3] === vector[3]
        @test_throws BoundsError dist[6]
        @test_throws BoundsError dist[3, 5]

        @test insupport(dist, vector)
        @test !insupport(dist, vector .+ 1e-6)
        @test !insupport(dist, vector .- 1e-6)

        @test pdf(dist, vector)         === 1.0
        @test pdf(dist, vector .+ 1e-6) === 0.0
        @test pdf(dist, vector .- 1e-6) === 0.0

        @test logpdf(dist, vector)         === 0.0
        @test logpdf(dist, vector .+ 1e-6) === -Inf
        @test logpdf(dist, vector .- 1e-6) === -Inf

        @test_throws MethodError insupport(dist, 1.0)
        @test_throws MethodError insupport(dist, [ 1.0 1.0; 1.0 1.0 ])
        @test_throws MethodError pdf(dist, 1.0)
        @test_throws MethodError pdf(dist, [ 1.0 1.0; 1.0 1.0 ])
        @test_throws MethodError logpdf(dist, 1.0)
        @test_throws MethodError logpdf(dist, [ 1.0 1.0; 1.0 1.0 ])

        @test entropy(dist) === -∞

        @test mean(dist)       == vector
        @test var(dist)        == zeros(5)
        @test std(dist)        == zeros(5)
        @test cov(dist)        == zeros(5, 5)
        @test precision(dist)  == fill(Inf, (5, 5))
        @test ndims(dist)      == 5

        @test probvec(dist)      == vector
        @test logmean(dist)      == log.(vector)
        @test_throws ErrorException inversemean(dist)
        @test_throws ErrorException mirroredlogmean(dist)
        @test loggammamean(dist) == loggamma.(vector)
    end

    @testset "Matrix-based PointMass" begin

        matrix = rand(5, 5)
        dist   = PointMass(matrix)

        @test variate_form(dist) === Matrixvariate
        @test dist[2]    === matrix[2]
        @test dist[3]    === matrix[3]
        @test dist[3, 3] === matrix[3, 3]
        @test_throws BoundsError dist[100]
        @test_throws BoundsError dist[100, 100]

        @test insupport(dist, matrix)
        @test !insupport(dist, matrix .+ 1e-6)
        @test !insupport(dist, matrix .- 1e-6)

        @test pdf(dist, matrix)         === 1.0
        @test pdf(dist, matrix .+ 1e-6) === 0.0
        @test pdf(dist, matrix .- 1e-6) === 0.0

        @test logpdf(dist, matrix)         === 0.0
        @test logpdf(dist, matrix .+ 1e-6) === -Inf
        @test logpdf(dist, matrix .- 1e-6) === -Inf

        @test_throws MethodError insupport(dist, 1.0)
        @test_throws MethodError insupport(dist, [ 1.0, 1.0 ])
        @test_throws MethodError pdf(dist, 1.0)
        @test_throws MethodError pdf(dist, [ 1.0, 1.0 ])
        @test_throws MethodError logpdf(dist, 1.0)
        @test_throws MethodError logpdf(dist, [ 1.0, 1.0 ])

        @test entropy(dist) === -∞

        @test mean(dist)       == matrix
        @test var(dist)        == zeros(5, 5)
        @test std(dist)        == zeros(5, 5)
        @test ndims(dist)      == (5, 5)

        @test_throws ErrorException cov(dist)
        @test_throws ErrorException precision(dist)

        @test_throws ErrorException probvec(dist)
        @test logmean(dist)      == log.(matrix)
        @test inversemean(dist)  ≈ cholinv(matrix)
        @test_throws ErrorException mirroredlogmean(dist)
        @test loggammamean(dist) == loggamma.(matrix)
    end

end

end


#Uitlity functions

" Spectral heat flux "
function transfer_w(T1,T2,w,τ)
    return (bose_einstein(w,T1)-bose_einstein(w,T2))*τ
end

"Total heat flux obtained from simulation"
function total_transfer(T1,T2,w,τ)
  t_w   = transfer_w.(T1,T2,w,τ)
  q_tot = trapz(w,t_w)
  return q_tot
end

"Generates frequency file that serves as imput to scuff-em"
function frequency_file(freqnum,T,typespace :: String ; f1=0.01 ,f2=10.0)
    freq = zeros(Float64,freqnum)
    if typespace == "linspace"
        freq =collect(linspace(f1*wien(T)/w0,f2*wien(T)/w0,freqnum))
        dffreq = DataFrame(Freq=freq)
        writetable("frequency_N="*string(freqnum)*"_"*typespace*".dat", dffreq,header=false)
    elseif typespace == "logspace"
        wi = log10(f1*wien(T)/w0)
        wf = log10(f2*wien(T)/w0)
        freq =collect(logspace(wi,wf,freqnum))
        dffreq = DataFrame(Freq=freq)
        writetable("frequency_N="*string(freqnum)*"_"*typespace*".dat", dffreq,header=false)
    end
end


#Function to return valid intersections of lines with image boundary

function get_valid_intersections(intersections::Vector{Tuple{T,U}}, indsx::AbstractUnitRange, indsy::AbstractUnitRange) where {T<:Real, U<:Real}
    valid_intersections = Vector{Tuple{T,U}}(undef, 0)
    for intersection in intersections
        if first(indsx) <= intersection[1] <= last(indsx) && first(indsy) <= intersection[2] <= last(indsy)
            push!(valid_intersections,intersection)
        end
    end
    valid_intersections
end

# LineTwoPoints methods

LineTwoPoints(x0::Int, y0::Int, x1::Int, y1::Int) = LineTwoPoints(Point(x0, y0), Point(x1,y1))
LineTwoPoints(p1::CartesianIndex{2}, p2::CartesianIndex{2}) = LineTwoPoints(Point(p1), Point(p2))

draw!(img::AbstractArray{T,2}, line::LineTwoPoints, method::Function = bresenham) where {T<:Colorant} =
    draw!(img, line, oneunit(T), method)

function draw!(img::AbstractArray{T,2}, line::LineTwoPoints, color::T, method::Function = bresenham) where T<:Colorant
    indsy, indsx = axes(img)
    x1 = line.p1.x; y1 = line.p1.y
    x2 = line.p2.x; y2 = line.p2.y
    m = (y2 - y1) / (x2 - x1)
    intersections_x = [(x, y1 + m*(x-x1)) for x in (first(indsx), last(indsx))]
    intersections_y = [(x1 + (y-y1)/m, y) for y in (first(indsy), last(indsy))]
    valid_intersections = get_valid_intersections(vcat(intersections_x, intersections_y), indsx, indsy)
    if length(valid_intersections) > 0
        method(img, round(Int,valid_intersections[1][2]), round(Int,valid_intersections[1][1]), round(Int,valid_intersections[2][2]), round(Int,valid_intersections[2][1]), color)
    else
        img
    end
end

# LineNormal methods

LineNormal(τ::Tuple{T,U}) where {T<:Real, U<:Real} = LineNormal(τ...)

draw!(img::AbstractArray{T,2}, line::LineNormal, method::Function = bresenham) where {T<:Colorant} =
    draw!(img, line, oneunit(T), method)

function draw!(img::AbstractArray{T, 2}, line::LineNormal, color::T, method::Function = bresenham) where T<:Colorant
    indsy, indsx = axes(img)
    cosθ = cos(line.θ)
    sinθ = sin(line.θ)
    intersections_x = [(x, (line.ρ - x*cosθ)/sinθ) for x in (first(indsx), last(indsx))]
    intersections_y = [((line.ρ - y*sinθ)/cosθ, y) for y in (first(indsy), last(indsy))]
    valid_intersections = get_valid_intersections(vcat(intersections_x, intersections_y), indsx, indsy)
    if length(valid_intersections) > 0
        method(img, round(Int,valid_intersections[1][2]), round(Int,valid_intersections[1][1]), round(Int,valid_intersections[2][2]), round(Int,valid_intersections[2][1]), color)
    else
        img
    end
end

# LineSegment methods

LineSegment(x0::Int, y0::Int, x1::Int, y1::Int) = LineSegment(Point(x0, y0), Point(x1,y1))
LineSegment(p1::CartesianIndex, p2::CartesianIndex) = LineSegment(Point(p1), Point(p2))

draw!(img::AbstractArray{T,2}, line::LineSegment, method::Function = bresenham) where {T<:Colorant} =
    draw!(img, line, oneunit(T), method)

draw!(img::AbstractArray{T,2}, line::LineSegment, color::T, method::Function = bresenham) where {T<:Colorant} =
    method(img, line.p1.y, line.p1.x, line.p2.y, line.p2.x, color)

# Methods to draw lines

function bresenham(img::AbstractArray{T, 2}, y0::Int, x0::Int, y1::Int, x1::Int, color::T) where T<:Colorant
    dx = abs(x1 - x0)
    dy = abs(y1 - y0)

    sx = x0 < x1 ? 1 : -1
    sy = y0 < y1 ? 1 : -1;

    err = (dx > dy ? dx : -dy) / 2

    while true
        drawifinbounds!(img, y0, x0, color)
        (x0 != x1 || y0 != y1) || break
        e2 = err
        if e2 > -dx
            err -= dy
            x0 += sx
        end
        if e2 < dy
            err += dx
            y0 += sy
        end
    end

    img
end

fpart(pixel::T) where {T} = pixel - T(trunc(pixel))
rfpart(pixel::T) where {T} = oneunit(T) - fpart(pixel)

function swap(x, y)
    y, x
end

function xiaolin_wu(img::AbstractArray{T, 2}, y0::Int, x0::Int, y1::Int, x1::Int, color::T) where T<:Gray
    dx = x1 - x0
    dy = y1 - y0

    swapped=false
    if abs(dx) < abs(dy)
        x0, y0 = swap(x0, y0)
        x1, y1 = swap(x1, y1)
        dx, dy = swap(dx, dy)
        swapped=true
    end
    if x1 < x0
        x0, x1 = swap(x0, x1)
        y0, y1 = swap(y0, y1)
    end
    gradient = dy / dx

    xend = round(Int, x0)
    yend = y0 + gradient * (xend - x0)
    xgap = rfpart(x0 + 0.5)
    xpxl0 = xend
    ypxl0 = trunc(Int, yend)
    index = swapped ? CartesianIndex(xpxl0, ypxl0) : CartesianIndex(ypxl0, xpxl0)
    drawifinbounds!(img, index, T(rfpart(yend) * xgap))
    index = swapped ? CartesianIndex(xpxl0, ypxl0 + 1) : CartesianIndex(ypxl0 + 1, xpxl0)
    drawifinbounds!(img, index, T(fpart(yend) * xgap))
    intery = yend + gradient

    xend = round(Int, x1)
    yend = y1 + gradient * (xend - x1)
    xgap = fpart(x1 + 0.5)
    xpxl1 = xend
    ypxl1 = trunc(Int, yend)
    index = swapped ? CartesianIndex(xpxl1, ypxl1) : CartesianIndex(ypxl1, xpxl1)
    drawifinbounds!(img, index, T(rfpart(yend) * xgap))
    index = swapped ? CartesianIndex(xpxl1, ypxl1 + 1) : CartesianIndex(ypxl1 + 1, xpxl1)
    drawifinbounds!(img, index, T(fpart(yend) * xgap))

    for i in (xpxl0 + 1):(xpxl1 - 1)
        index = swapped ? CartesianIndex(i, trunc(Int, intery)) : CartesianIndex(trunc(Int, intery), i)
        drawifinbounds!(img, index, T(rfpart(intery)))
        index = swapped ? CartesianIndex(i, trunc(Int, intery) + 1) : CartesianIndex(trunc(Int, intery) + 1, i)
        drawifinbounds!(img, index, T(fpart(intery)))
        intery += gradient
    end
    img
end


using Test

@testset "Broyden ($R)" for R in [Float32, Float64]
    using LinearAlgebra
    using ProximalAlgorithms: Broyden, update!

    H = R[
        0.63287    0.330934   -0.156908   -0.294776    0.10761;
        0.330934   0.673201    0.0459778   0.231011   -0.235265;
       -0.156908   0.0459778   0.635812   -0.232261   -0.388775;
       -0.294776   0.231011   -0.232261    0.726854   -0.0691783;
        0.10761   -0.235265   -0.388775   -0.0691783   0.336262;
    ]
    l = R[1.0, 2.0, 3.0, 4.0, 5.0]

    n = length(l)

    f(x) = dot(x, H * x) / 2 + dot(x, l)

    x_star = -H \ l
    f_star = f(x_star)
    norm_x_star = norm(x_star)

    Lip = opnorm(H)
    gamma = 1 / Lip
    x = zeros(n)
    res = zeros(n)
    x_prev = zeros(n)
    res_prev = zeros(n)

    acc = Broyden(x)

    mul!(res, H, x)
    res .+= l
    res .*= gamma

    d = acc * res

    for it = 1:10
        # store iterate and residual for the operator update later
        res_prev .= res
        x_prev .= x

        # compute accelerated direction
        mul!(d, acc, res)

        # update iterate
        x .-= d

        # compute new residual
        mul!(res, H, x)
        res .+= l
        res .*= gamma

        # update operator
        update!(acc, x - x_prev, res - res_prev)
    end

    @test f(x) <= f_star + (1 + abs(f_star)) * sqrt(eps(R))

end


module ForcedFlowFreeSlip

using Printf

using Oceananigans
using Oceananigans.BoundaryConditions
using Oceananigans.OutputWriters
using Oceananigans.Fields

# Functions that define the forced flow problem

 ξ(t) = 1 + sin(t^2)
ξ′(t) = 2t * cos(t^2)

 f(x, t) =   cos(x - ξ(t))
fₓ(x, t) = - sin(x - ξ(t))

Fᵘ(x, y, t) = (4 * f(x, t) - 2 * ξ′(t) * fₓ(x, t)) * cos(y)
Fᵛ(x, y, t) = sin(2y) / 2

u(x, y, t) =   f(x, t) * cos(y)
v(x, y, t) = -fₓ(x, t) * sin(y)

#####
##### x, z
#####

function print_progress(simulation)
    model = simulation.model
    i, t = model.clock.iteration, model.clock.time
    progress = 100 * (i / simulation.stop_iteration)
    @info @sprintf("[%05.2f%%] iteration: %d, time: %.5e", progress, i, t)
    return nothing
end

const DATA_DIR = joinpath(@__DIR__, "..", "data")

function setup_xz_simulation(; Nx, Δt, stop_iteration, architecture=CPU(), dir=DATA_DIR)

    grid = RegularRectilinearGrid(size=(Nx, 1, Nx), x=(0, 2π), y=(0, 1), z=(0, π),
                                topology=(Periodic, Periodic, Bounded))

    model = IncompressibleModel(architecture = architecture,
                                        grid = grid,
                                    coriolis = nothing,
                                    buoyancy = nothing,
                                     tracers = nothing,
                                     closure = IsotropicDiffusivity(ν=1),
                                     forcing = (u = (x, y, z, t) -> Fᵘ(x, z, t),
                                                w = (x, y, z, t) -> Fᵛ(x, z, t))
                                )

    set!(model, u = (x, y, z) -> u(x, z, 0),
                w = (x, y, z) -> v(x, z, 0))

    simulation = Simulation(model, Δt=Δt, stop_iteration=stop_iteration, progress=print_progress, iteration_interval=40)

    simulation.output_writers[:fields] = JLD2OutputWriter(model, model.velocities;
                                                          dir = dir, force = true,
                                                          field_slicer = nothing,
                                                          prefix = @sprintf("forced_free_slip_xz_Nx%d_Δt%.1e", Nx, Δt),
                                                          schedule = TimeInterval(stop_iteration * Δt / 2))

    return simulation
end

function setup_and_run_xz(; setup...)
    simulation = setup_xz_simulation(; setup...)
    @info "Running forced flow free slip simulation with Nx = Nz = $(setup[:Nx]), Δt = $(setup[:Δt])"
    @time run!(simulation)
    return nothing
end

#####
##### x, y
#####

function setup_xy_simulation(; Nx, Δt, stop_iteration, architecture=CPU(), dir=DATA_DIR)

    grid = RegularRectilinearGrid(size=(Nx, Nx, 1), x=(0, 2π), y=(0, π), z=(0, 1),
                                topology=(Periodic, Bounded, Bounded))

    model = IncompressibleModel(architecture = architecture,
                                        grid = grid,
                                    coriolis = nothing,
                                    buoyancy = nothing,
                                     tracers = nothing,
                                     closure = IsotropicDiffusivity(ν=1),
                                     forcing = (u = (x, y, z, t) -> Fᵘ(x, y, t),
                                                v = (x, y, z, t) -> Fᵛ(x, y, t))
                                )

    set!(model, u = (x, y, z) -> u(x, y, 0),
                v = (x, y, z) -> v(x, y, 0))

    simulation = Simulation(model, Δt=Δt, stop_iteration=stop_iteration, progress=print_progress, iteration_interval=40)

    simulation.output_writers[:fields] = JLD2OutputWriter(model, model.velocities;
                                                          dir = dir, force = true,
                                                          field_slicer = nothing,
                                                          prefix = @sprintf("forced_free_slip_xy_Nx%d_Δt%.1e", Nx, Δt),
                                                          schedule = TimeInterval(stop_iteration * Δt / 2))

    return simulation
end

function setup_and_run_xy(; setup...)
    simulation = setup_xy_simulation(; setup...)
    @info "Running forced flow free slip simulation with Nx = Ny = $(setup[:Nx]), Δt = $(setup[:Δt])"
    @time run!(simulation)
    return nothing
end

end # module


import Singular.PolynomialRing
import Singular.QQ

#-------------------------------------------------------------------------------

funcname(x) = PyCall.PyObject(x).func.__name__

fn_map = Dict(
              "Add" => :+,
              "Sub" => :-,
              "Mul" => :*,
              "Div" => :/,
              "Pow" => :^
)

val_map = Dict(
               "Zero"             => (0),
               "One"              => (1),
               "NegativeOne"      => (-1),
               "Half"             => QQ(1)//QQ(2),
            #    "Pi"               => :pi,
            #    "Exp1"             => :e,
            #    "Infinity"         => :Inf,
            #    "NegativeInfinity" => :(-Inf),
            #    "ComplexInfinity"  => :Inf, # error?
            #    "ImaginaryUnit"    => :im,
            #    "BooleanTrue"      => :true,
            #    "BooleanFalse"     => :false
)

# map_fn(fn_map, key)   = get(fn_map, key, key)
# map_var(var_map, key) = get(var_map, key, key)

function replace_sympy(ex, vars::Dict)
    fn = funcname(ex)

    if fn == "Symbol"
        return get(vars, string(ex), ex)
    elseif fn in ["Integer" , "Float"]
        return N(ex)
    elseif fn == "Rational"
        return QQ(convert(Int, numer(ex)))//QQ(convert(Int, denom(ex)))
    elseif haskey(val_map, fn)
        return val_map[fn]
    end
    error("should not be visited")
    # Expr(:call, map_fn(fns_map, fn), [replace_sympy(a, values=values) for a in args(ex)]...)
end

@syms x y
const ADD = (x+y).func
const MUL = (x*y).func
const POW = (x^y).func

function sym2spoly(b::Array{Sym, 1})
    fvars = union(free_symbols.(b)...)
    svars = string.(fvars)
    R, pvars = PolynomialRing(QQ, svars)
    dict = Dict(zip(string.(fvars), pvars))
    basis = [shallow_replace(p, dict) for p in b]
    basis = eval.(basis)
    return R, basis, dict
end

function shallow_replace(expr::Sym, vars::Dict)
    fn = funcname(expr)
    # println("Function: ", fn==ADD)
    a = [shallow_replace(arg, vars) for arg in expr.args]
    # println("Start: ", expr)
    # println("Start: ", a)
    if fn == "Add"
        return sum(a)
    elseif fn == "Mul"
        return prod(a)
    elseif fn == "Pow"
        return a[1]^a[2]
    end
    return replace_sympy(expr, vars)
end

# function sym2spoly(p::Sym; vars = Sym[])
#     sym2spoly([p], vars=vars)[1]
# end

function sym2spoly(poly::Sym, varmap::Dict)
    # fns_map = merge(fn_map, fns)
    # vals_map = merge(val_map, varmap)
    eval(shallow_replace(poly, varmap))
end

#-------------------------------------------------------------------------------

# function Ideal(basis::Array{Sym, 1}; vars = Sym[])
#     R, sbasis, varmap = sym2spoly(basis, vars)
#     return Singular.Ideal(R, sbasis...), varmap
# end

function Ideal(basis::Array{Sym, 1})
    R, sbasis, varmap = sym2spoly(basis)
    return Singular.Ideal(R, sbasis...), varmap
end

groebner(ideal::Singular.sideal) = std(ideal)

function eliminate(basis::Array{Sym, 1}, elim::Array{Sym, 1})
    ideal, varmap = Ideal(basis)
    vars = [varmap[string(v)] for v in elim]
    return Singular.eliminate(ideal, vars...)
end


""" Computing in the category of finite sets and functions.
"""
module FinSets
export FinSet, FinFunction, FinFunctionCallable, FinFunctionVector, force

using Compat: only

using AutoHashEquals
using DataStructures: IntDisjointSets, union!, find_root
import FunctionWrappers: FunctionWrapper

using ...GAT
using ...Theories: Category
import ...Theories: dom, codom, id, compose, ⋅, ∘
using ..FreeDiagrams, ..Limits
import ..Limits: limit, colimit, universal

# Data types
############

""" Abstract type for objects in a category that are sets.

An implementation detail used to control method dispatch. Concrete subtypes
include `FinSet` and `FinRel`.
"""
abstract type AbstractSetOb{S,T} end

Base.eltype(::Type{AbstractSetOb{S}}) where S = eltype(S)
Base.iterate(s::AbstractSetOb, args...) = iterate(iterable(s), args...)
Base.length(s::AbstractSetOb) = length(iterable(s))
Base.in(s::AbstractSetOb, elem) = in(s, iterable(s))

""" Finite set.

This generic type encompasses the category **FinSet** of finite sets and
functions, through types `FinSet{S} where S <: AbstractSet`, as well as the
skeleton of this category, through the type `FinSet{Int}`. In the latter case,
the object `FinSet(n)` represents the set {1,...,n}.
"""
@auto_hash_equals struct FinSet{S,T} <: AbstractSetOb{S,T}
  set::S
end

FinSet(n::Int) = FinSet{Int,Int}(n)
FinSet(set::S) where {T,S<:AbstractSet{T}} = FinSet{S,T}(set)
iterable(s::FinSet{Int}) = 1:s.set
iterable(s::FinSet{<:AbstractSet}) = s.set

Base.show(io::IO, s::FinSet) = print(io, "FinSet($(s.set))")

""" Function between finite sets.

The function can be defined implicitly by an arbitrary Julia function, in which
case it is evaluated lazily, or explictly by a vector of integers. In the latter
case, the function (1↦1, 2↦3, 3↦2, 4↦3), for example, is represented by the
vector [1,3,2,3].
"""
abstract type FinFunction{S,T} end

FinFunction(f::Function, args...) = FinFunctionCallable(f, args...)
FinFunction(f::AbstractVector, args...) = FinFunctionVector(f, args...)
FinFunction(::typeof(identity), args...) = FinFunctionIdentity(args...)

""" Function in FinSet defined by a callable Julia object.

To be evaluated lazily unless forced.
"""
@auto_hash_equals struct FinFunctionCallable{S,T} <: FinFunction{S,T}
  func::FunctionWrapper{T,Tuple{T}} # Usually `Function` but can be any Julia callable.
  dom::FinSet{S,T}
  codom::FinSet{S,T}
end
FinFunctionCallable(f::Function, dom::Int, codom::Int) =
  FinFunctionCallable(FunctionWrapper{Int,Tuple{Int}}(f), FinSet(dom), FinSet(codom))
FinFunctionCallable(f::Function, dom::FinSet{S,T}, codom::FinSet{S,T}) where {S,T} =
  FinFunctionCallable(FunctionWrapper{T,Tuple{T}}(f), dom, codom)

function Base.show(io::IO, f::FinFunctionCallable)
  func = f.func.obj[] # Deference FunctionWrapper
  print(io, "FinFunction($(nameof(func)), $(f.dom), $(f.codom))")
end

(f::FinFunctionCallable)(x) = f.func(x)

""" Function in FinSet represented explicitly by a vector.

The elements of the set are assumed to be {1,...,n}.
"""
@auto_hash_equals struct FinFunctionVector{T<:AbstractVector{Int}} <: FinFunction{Int,Int}
  func::T
  codom::Int
end
FinFunctionVector(f::AbstractVector) =
  FinFunctionVector(f, isempty(f) ? 0 : maximum(f))

function FinFunctionVector(f::AbstractVector, dom::Int, codom::Int)
  length(f) == dom || error("Length of vector $f does not match domain $dom")
  FinFunctionVector(f, codom)
end

FinFunctionVector(f::AbstractVector, codom::FinSet{Int}) =
  FinFunctionVector(f, length(codom))
FinFunctionVector(f::AbstractVector, dom::FinSet{Int}, codom::FinSet{Int}) =
  FinFunctionVector(f, length(dom), length(codom))

dom(f::FinFunctionVector) = FinSet(length(f.func))
codom(f::FinFunctionVector) = FinSet(f.codom)

Base.show(io::IO, f::FinFunctionVector) =
  print(io, "FinFunction($(f.func), $(length(f.func)), $(f.codom))")

(f::FinFunctionVector)(x) = f.func[x]

""" Force evaluation of lazy function or relation.
"""
force(f::FinFunction{Int}) = FinFunctionVector(map(f, dom(f)), codom(f))
force(f::FinFunctionVector) = f

Base.collect(f::FinFunction) = force(f).func

""" Identity function in FinSet.
"""
@auto_hash_equals struct FinFunctionIdentity{S,T} <: FinFunction{S,T}
  dom::FinSet{S,T}
end

function FinFunctionIdentity(dom, codom)
  @assert dom == codom
  FinFunctionIdentity(dom)
end
FinFunctionIdentity(n::Int) = FinFunctionIdentity(FinSet(n))

dom(f::FinFunctionIdentity) = f.dom
codom(f::FinFunctionIdentity) = f.dom

Base.show(io::IO, f::FinFunctionIdentity) =
  print(io, "FinFunction(identity, $(f.dom))")

(f::FinFunctionIdentity)(x) = x

# Category of finite sets
#########################

""" Category of finite sets and functions.
"""
@instance Category{FinSet, FinFunction} begin
  dom(f::FinFunction) = f.dom
  codom(f::FinFunction) = f.codom

  id(A::FinSet) = FinFunction(identity, A, A)

  function compose(f::FinFunction, g::FinFunction)
    @assert codom(f) == dom(g)
    compose_impl(f, g)
  end
end

compose_impl(f::FinFunction{S,T}, g::FinFunction{S,T}) where {S,T} =
  FinFunction(g ∘ f, dom(f), codom(g))
compose_impl(f::FinFunctionVector, g::FinFunctionVector) =
  FinFunctionVector(g.func[f.func], g.codom)

compose_impl(f::FinFunction, ::FinFunctionIdentity) = f
compose_impl(::FinFunctionIdentity, f::FinFunction) = f
compose_impl(f::FinFunctionIdentity, ::FinFunctionIdentity) = f

# Limits
########

function limit(Xs::EmptyDiagram{<:FinSet{Int}})
  Limit(Xs, SMultispan{0}(FinSet(1)))
end

function universal(lim::Terminal{<:FinSet{Int}},
                   cone::SMultispan{0,<:FinSet{Int}})
  FinFunction(ones(Int, length(apex(cone))))
end

function limit(Xs::ObjectPair{<:FinSet{Int}})
  m, n = length.(Xs)
  indices = CartesianIndices((m, n))
  π1 = FinFunction(i -> indices[i][1], m*n, m)
  π2 = FinFunction(i -> indices[i][2], m*n, n)
  Limit(Xs, Span(π1, π2))
end

function universal(lim::BinaryProduct{<:FinSet{Int}}, cone::Span{<:FinSet{Int}})
  f, g = cone
  m, n = length.(codom.(cone))
  indices = LinearIndices((m, n))
  FinFunction(i -> indices[f(i),g(i)], apex(cone), ob(lim))
end

function limit(Xs::DiscreteDiagram{<:FinSet{Int}})
  ns = length.(Xs)
  indices = CartesianIndices(Tuple(ns))
  n = prod(ns)
  πs = [FinFunction(i -> indices[i][j], n, ns[j]) for j in 1:length(ns)]
  Limit(Xs, Multispan(FinSet(n), πs))
end

function universal(lim::Product{<:FinSet{Int}}, cone::Multispan{<:FinSet})
  ns = length.(codom.(cone))
  indices = LinearIndices(Tuple(ns))
  FinFunction(i -> indices[(f(i) for f in cone)...], apex(cone), ob(lim))
end

function limit(pair::ParallelPair{<:FinSet{Int}})
  f, g = pair
  m = length(dom(pair))
  eq = FinFunction(filter(i -> f(i) == g(i), 1:m), m)
  Limit(pair, SMultispan{1}(eq))
end

function limit(para::ParallelMorphisms{<:FinSet{Int}})
  @assert !isempty(para)
  f1, frest = para[1], para[2:end]
  m = length(dom(para))
  eq = FinFunction(filter(i -> all(f1(i) == f(i) for f in frest), 1:m), m)
  Limit(para, SMultispan{1}(eq))
end

function universal(lim::Equalizer{<:FinSet{Int}},
                   cone::SMultispan{1,<:FinSet{Int}})
  ι = collect(incl(lim))
  h = only(cone)
  FinFunction(Int[only(searchsorted(ι, i)) for i in collect(h)], length(ι))
end

limit(cospan::Multicospan{<:FinSet{Int}}) = composite_pullback(cospan)

function limit(d::FreeDiagram{<:FinSet{Int}})
  p = product(ob(d))
  n, leg = length(ob(p)), legs(p)
  satisfy(e,x) = hom(d,e)(leg[src(d,e)](x)) == leg[tgt(d,e)](x)
  f = FinFunction(filter(i -> all(satisfy(e,i) for e in edges(d)), 1:n), n)
  Limit(d, Multispan(dom(f), [compose(f,leg[i]) for i in vertices(d)]))
end

# Colimits
##########

function colimit(Xs::EmptyDiagram{<:FinSet{Int}})
  Colimit(Xs, SMulticospan{0}(FinSet(0)))
end

function universal(colim::Initial{<:FinSet{Int}},
                   cocone::SMulticospan{0,<:FinSet{Int}})
  FinFunction(Int[], apex(cocone))
end

function colimit(Xs::ObjectPair{<:FinSet{Int}})
  m, n = length.(Xs)
  ι1 = FinFunction(1:m, m, m+n)
  ι2 = FinFunction(m+1:m+n, n, m+n)
  Colimit(Xs, Cospan(ι1, ι2))
end

function universal(colim::BinaryCoproduct{<:FinSet{Int}},
                   cocone::Cospan{<:FinSet{Int}})
  f, g = cocone
  FinFunction(vcat(collect(f), collect(g)), ob(colim), apex(cocone))
end

function colimit(Xs::DiscreteDiagram{<:FinSet{Int}})
  ns = length.(Xs)
  n = sum(ns)
  offsets = [0,cumsum(ns)...]
  ιs = [FinFunction((1:ns[j]) .+ offsets[j],ns[j],n) for j in 1:length(ns)]
  Colimit(Xs, Multicospan(FinSet(n), ιs))
end

function universal(colim::Coproduct{<:FinSet{Int}},
                   cocone::Multicospan{<:FinSet{Int}})
  FinFunction(reduce(vcat, (collect(f) for f in cocone), init=Int[]),
              ob(colim), apex(cocone))
end

function colimit(pair::ParallelPair{<:FinSet{Int}})
  f, g = pair
  m, n = length(dom(pair)), length(codom(pair))
  sets = IntDisjointSets(n)
  for i in 1:m
    union!(sets, f(i), g(i))
  end
  h = [ find_root(sets, i) for i in 1:n ]
  roots = unique!(sort(h))
  coeq = FinFunction([ searchsortedfirst(roots, r) for r in h], length(roots))
  Colimit(pair, SMulticospan{1}(coeq))
end

function colimit(para::ParallelMorphisms{<:FinSet{Int}})
  @assert !isempty(para)
  f1, frest = para[1], para[2:end]
  m, n = length(dom(para)), length(codom(para))
  sets = IntDisjointSets(n)
  for i in 1:m
    for f in frest
      union!(sets, f1(i), f(i))
    end
  end
  h = [ find_root(sets, i) for i in 1:n ]
  roots = unique!(sort(h))
  coeq = FinFunction([ searchsortedfirst(roots, r) for r in h ], length(roots))
  Colimit(para, SMulticospan{1}(coeq))
end

function universal(coeq::Coequalizer{<:FinSet{Int}},
                   cocone::SMulticospan{1,<:FinSet{Int}})
  q = zeros(Int, length(ob(coeq)))
  π = proj(coeq)
  h = only(cocone)
  for i in dom(h)
    j = π(i)
    if q[j] == 0
      q[j] = h(i)
    else
      @assert q[j] == h(i) "Quotient map out of coequalizer is ill-defined"
    end
  end
  FinFunction(q, codom(h))
end

colimit(span::Multispan{<:FinSet{Int}}) = composite_pushout(span)

function colimit(d::FreeDiagram{<:FinSet{Int}})
  coprod = coproduct(ob(d))
  n, leg = length(ob(coprod)), legs(coprod)
  sets = IntDisjointSets(n)
  for e in edges(d)
    s, t, h = src(d,e), tgt(d,e), hom(d,e)
    for i in dom(h)
      union!(sets, leg[s](i), leg[t](h(i)))
    end
  end
  h = [ find_root(sets, i) for i in 1:n ]
  roots = unique!(sort(h))
  m = length(roots)
  f = FinFunction([ searchsortedfirst(roots, r) for r in h ], m)
  Colimit(d, Multicospan(FinSet(m), [ compose(leg[i],f) for i in vertices(d) ]))
end

end


import Optim, LineSearches

struct GradientAscent <: Algorithm end

function train(learner::Learner{GradientAscent})
    num_samples = sample_count(learner.train_data)
    β_init = _unif_random(learner.param_lb, learner.param_ub)
    offset_init = rand(Uniform(learner.offset_lb, learner.offset_ub))

    l1_coeff = 0.0
    l2_coeff = 0.0
    for penalty in learner.objective_penalties
        if penalty isa L1Regularizer
            @assert l1_coeff == 0.0
            l1_coeff = penalty.coefficient
        elseif penalty isa L2Regularizer
            @assert l2_coeff = 0.0
            l2_coeff = penalty.coefficient
        else
            error("Unrecognized type of regularizer")
        end
    end
    function l(x)
        β = x[1:(end-1)]
        offset = x[end]
        model = LinearModel(β, offset)
        return -1 * (1 / num_samples * sum(_reward_function(data_point, model) for data_point in learner.train_data.samples) +
            l1_coeff * sum(abs.(β)) +
            l2_coeff * sum(β.^2))
    end
    function ∇l(x)
        β = x[1:(end-1)]
        offset = x[end]
        model = LinearModel(β, offset)
        return -1 * (1 / num_samples * sum(vcat(_grad_reward_function(data_point, model)...) for data_point in learner.train_data.samples) +
            l1_coeff * vcat(sign.(β), 0.0) +
            2l2_coeff * vcat(β, 0.0))
    end

    results = Optim.optimize(l, ∇l, vcat(β_init, offset_init), Optim.GradientDescent(linesearch=LineSearches.StrongWolfe()), Optim.Options(show_trace=false, allow_f_increases=true, extended_trace=false), inplace=false)

    β_values = results.minimizer[1:(end-1)]
    offset_value = results.minimizer[end]

    return TrainingResults(learner, LinearModel(β_values, offset_value))
end


# Download BedMachine
using Pkg
Pkg.activate(".")
using NCDatasets, Plots, Parameters

!ispath("../data") && mkdir("../data")
!ispath("../output") && mkdir("../output")

# smooth_surface = 2000.0 # smoothing radius, in meters
nsm_s  = 2 # surface
nsm_b  = 0 # bed

run_id = "run_ds_20" # DEBUG: define the output directory within ../output/

!ispath("../output/$run_id") && mkdir("../output/$run_id")

const rw = 1000.0
const ri = 910.0

url = "file://" * homedir() * "/itet-stor/glazioarch/GlacioData/BedMachine_Antarctica/168596330/BedMachineAntarctica_2019-11-05_v01.nc"
# url = "file://" * "/Volumes/luraess/glazioarch/GlacioData/BedMachine_Antarctica/168596330/BedMachineAntarctica_2019-11-05_v01.nc"

filename = if !isfile(joinpath("../data", splitdir(url)[2]))
    print("Downloading data ...")
    download(url, joinpath("../data", splitdir(url)[2]))
    println("done.")
else
    joinpath("../data", splitdir(url)[2])
end

# if !(@isdefined ds) # DEBUG: should be done differently to allow for changing ds while the file already downloaded
    print("Loading NC file ... ")
    ds = NCDataset(filename)

    # downscale = 2 finest possible on 16GB
    # downscale = 1 ok on 32GB ram (needs around 16GB)
    downscale = 30

    x  = ds["x"][1:downscale:end]
    y  = ds["y"][1:downscale:end]
    x  = x[1]:x[2]-x[1]:x[end]
    y  = y[1]:y[2]-y[1]:y[end]
    mask_   = ds["mask"][1:downscale:end, 1:downscale:end]
    mask    = (mask_.==2) .| (mask_.==4) .| (mask_.==1) # route water over(under): grounded ice or lake Vostok or ice-free-land
    surface = ds["surface"][1:downscale:end, 1:downscale:end]
    bed     = convert(Matrix{Float32}, ds["bed"][1:downscale:end, 1:downscale:end])

    # Pad the data to make it GPU compatible with good performance # DEBUG: do this until ParallelStencil is fixed: see https://github.com/samo-lin/ParallelStencil.jl/issues/42
    olx = rem(length(x), 32)
    oly = rem(length(y), 8 )
    x1=1; xE=0; if (olx!=0) x1 = Int(ceil(olx/2)); xE = olx-x1+1; end
    y1=1; yE=0; if (oly!=0) y1 = Int(ceil(oly/2)); yE = oly-y1+1; end

    x  = x[x1:end-xE]
    y  = y[y1:end-yE]
    mask    = mask[x1:end-xE, y1:end-yE]
    surface = surface[x1:end-xE, y1:end-yE]
    bed     = bed[x1:end-xE, y1:end-yE]

    println("done")
# end
println("grid size: nx=$(length(x)), ny=$(length(y))")

# if smooth_surface>0
#     dn = round(Int, smooth_surface/step(x))
#     if dn>0
#         #smooth
#         @warn "todo"
#     end
# end

@views function smooth!(A)
    A[2:end-1,2:end-1] .= A[2:end-1,2:end-1] .+ 1.0./6.1.*(diff(diff(A[:,2:end-1],dims=1),dims=1) .+ diff(diff(A[2:end-1,:],dims=2),dims=2))
    return
end

print("Smoothing DEM [surface $nsm_s step(s), bed $nsm_b step(s)] ... ")
for ism = 1:nsm_s
    smooth!(surface)
end
for ism = 1:nsm_b
    smooth!(bed)
end
println("done")

surf = copy(surface)
surf[.!mask] .= NaN

## Now run flow routing with surface and bed

# using WhereTheWaterFlows, PyPlot
# const WWF = WhereTheWaterFlows
# inds = [45:55,45:55]
# mask = mask[inds...]
# hydpot = (rw*bed .+ (surf.-bed)*ri)[inds...]

# hydpot = (rw*bed .+ (surf.-bed)*ri)
# @time area, slen, dir, nout, nin, pits, c, bnds = waterflows( hydpot, drain_pits=true, bnd_as_pits=true);
# area[.!mask] .= NaN;
# WWF.plotarea(1:size(area,1), 1:size(area,2), area, pits)
# if !all(getindex.(Ref(mask), pits).==0)
#     @warn "Not all pits were removed"
# end
#

using WhereTheWaterAlsoFlows
#const WWAF = WhereTheWaterAlsoFlows

@time D, F, xcp, ycp = flow_routing2D(x, reverse(y), bed, surface.-bed, mask, mask*1, outdir="../output/$run_id")

# visu
F[.!mask] .= NaN

plt = heatmap(xcp, ycp, log10.(F'), aspect_ratio=1, xlims=(xcp[1], xcp[end]), ylims=(ycp[1], ycp[end]), c=:viridis, framestyle=:box, title="log10 ||Flux||")

# savefig(plot(plt, dpi=300), joinpath(@__DIR__, "../output/$run_id/o$(length(x))_final.png"))

display(plt)

nothing