"""
    ReLU(θ, γ)

ReLU layer, with location parameters `θ` and scale parameters `γ`.
"""
struct ReLU{N,A1,A2} <: AbstractLayer{N}
    θ::A1
    γ::A2
    function ReLU(θ::AbstractArray, γ::AbstractArray)
        @assert size(θ) == size(γ)
        return new{ndims(θ), typeof(θ), typeof(γ)}(θ, γ)
    end
end

ReLU(::Type{T}, sz::Int...) where {T} = ReLU(zeros(T, sz...), ones(T, sz...))
ReLU(sz::Int...) = ReLU(Float64, sz...)

Base.repeat(l::ReLU, n::Int...) = ReLU(repeat(l.θ, n...), repeat(l.γ, n...))

function energies(layer::ReLU, x::AbstractArray)
    @assert size(layer) == size(x)[1:ndims(layer)]
    return relu_energy.(layer.θ, layer.γ, x)
end

free_energies(layer::ReLU, inputs::Union{Real,AbstractArray} = 0) = relu_free.(layer.θ .+ inputs, layer.γ)
sample_from_inputs(layer::ReLU, inputs::Union{Real,AbstractArray} = 0) = relu_rand.(layer.θ .+ inputs, layer.γ)
mode_from_inputs(layer::ReLU, inputs::Union{Real,AbstractArray} = 0) = max.((layer.θ .+ inputs) ./ abs.(layer.γ), 0)

function mean_from_inputs(layer::ReLU, inputs::Union{Real,AbstractArray} = 0)
    g = Gaussian(layer.θ, layer.γ)
    μ = mean_from_inputs(g, inputs)
    σ = sqrt.(var_from_inputs(g, inputs))
    return @. μ + σ * tnmean(-μ / σ)
end

mean_abs_from_inputs(layer::ReLU, inputs::Union{Real,AbstractArray} = 0) = mean_from_inputs(layer, inputs)

function var_from_inputs(layer::ReLU, inputs::Union{Real,AbstractArray} = 0)
    g = Gaussian(layer.θ, layer.γ)
    μ = mean_from_inputs(g, inputs)
    ν = var_from_inputs(g, inputs)
    return @. ν * tnvar(-μ / √ν)
end

function meanvar_from_inputs(layer::ReLU, inputs::Union{Real,AbstractArray} = 0)
    g = Gaussian(layer.θ, layer.γ)
    μ = mean_from_inputs(g, inputs)
    ν = var_from_inputs(g, inputs)
    σ = sqrt.(ν)
    tμ, tν = tnmeanvar(-μ ./ σ)
    return μ + σ .* tμ, ν .* tν
end

std_from_inputs(layer::ReLU, inputs::Union{Real,AbstractArray} = 0) = sqrt.(var_from_inputs(layer, inputs))

function ∂free_energies(layer::ReLU, inputs::Union{Real,AbstractArray} = 0)
    μ, ν = meanvar_from_inputs(layer, inputs)
    return (θ = -μ, γ = sign.(layer.γ) .* (ν .+ μ.^2) / 2)
end

const ReLUStats = GaussStats
suffstats(layer::ReLU, data::AbstractArray; wts = nothing) = ReLUStats(layer, data; wts)

function ∂energy(layer::ReLU, stats::ReLUStats)
    @assert size(layer) == size(stats)
    return (; θ = -stats.x1, γ = sign.(layer.γ) .* stats.x2 / 2)
end

function relu_energy(θ::Real, γ::Real, x::Real)
    E = gauss_energy(θ, γ, x)
    return x < 0 ? inf(E) : E
end

function relu_free(θ::Real, γ::Real)
    abs_γ = abs(γ)
    return -logerfcx(-θ / √(2abs_γ)) + log(2abs_γ/π) / 2
end

function relu_rand(θ::Real, γ::Real)
    abs_γ = abs(γ)
    μ = θ / abs_γ
    σ = √inv(abs_γ)
    return randnt_half(μ, σ)
end


"""
# References
[1] M. Islam, M. Okasha, and M. M. Idres, “Dynamics and control of quadcopter using linear model predictive control approach,” IOP Conf. Ser.: Mater. Sci. Eng., vol. 270, p. 012007, Dec. 2017, doi: 10.1088/1757-899X/270/1/012007.
# Variables
u ∈ R^4: square of rotor angular speed
"""
Base.@kwdef struct IslamQuadcopter <: Quadcopter
    J = Diagonal([7.5e-3, 7.5e-3, 1.3e-2])  # kg m^2
    l = 0.23  # m
    kf = 3.13e-5  # N s^2
    kM = 7.5e-7  # N m s^2
    m = 0.65  # kg
    g = 9.81  # m / s^2
    B = [kf  kf  kf  kf;
          0 -kf   0  kf;
         kf   0 -kf   0;
         kM -kM  kM -kM]
    # drags
    # kt = Diagonal(0.1*ones(3))  # Ns / m
    # kr = Diagonal(0.1*ones(3))  # N m s
    # input limits
    dim_input = 4
    u_min = zeros(dim_input)
    u_max = (m*g/kf) * ones(dim_input)
end

function saturate(env::IslamQuadcopter, u)
    @unpack u_min, u_max = env
    u_saturated = clamp.(u, u_min, u_max)
end

function input_to_force_moment(env::IslamQuadcopter, u)
    @unpack B = env
    ν = B * u
end


export LSTM

mutable struct LSTM
    insize::Int
    hsize::Int
    nlayers::Int
    droprate::Float64
    bidirectional::Bool
    params::Vector          # W, U, b, h, c
end

function getparams(lstm::LSTM)
    params = Var[]
    for p in lstm.params
        append!(params, p)
    end
    params
end

doc"""
    LSTM(::Type{T}, insize::Int, hsize::Int, nlayers::Int, droprate::Float64, bidirectional::Bool,
        [init_w=Orthogonal(), init_u=Orthogonal(), init_b=Fill(0)])

Long Short-Term Memory network.

```math
\begin{align*}
\mathbf{f}_{t} & =\sigma_{g}(W_{f}\mathbf{x}_{t}+U_{f}\mathbf{h}_{t-1}+\mathbf{b}_{f})\\
\mathbf{i}_{t} & =\sigma_{g}(W_{i}\mathbf{x}_{t}+U_{i}\mathbf{h}_{t-1}+\mathbf{b}_{i})\\
\mathbf{o}_{t} & =\sigma_{g}(W_{o}\mathbf{x}_{t}+U_{o}\mathbf{h}_{t-1}+\mathbf{b}_{o})\\
\mathbf{c}_{t} & =\mathbf{f}_{t}\odot\mathbf{c}_{t-1}+\mathbf{i}_{t}\odot\sigma_{c}(W_{c}\mathbf{x}_{t}+U_{c}\mathbf{h}_{t-1}+\mathbf{b}_{c})\\
\mathbf{h}_{t} & =\mathbf{o}_{t}\odot\sigma_{h}(\mathbf{c}_{t})
\end{align*}
```

* ``x_t \in R^{d}``: input vector to the LSTM block
* ``f_t \in R^{h}``: forget gate's activation vector
* ``i_t \in R^{h}``: input gate's activation vector
* ``o_t \in R^{h}``: output gate's activation vector
* ``h_t \in R^{h}``: output vector of the LSTM block
* ``c_t \in R^{h}``: cell state vector
* ``W \in R^{h \times d}``, ``U \in R^{h \times h}`` and ``b \in R^{h}``: weight matrices and bias vectors
* ``\sigma_g``: sigmoid function
* ``\sigma_c``: hyperbolic tangent function
* ``\sigma_h``: hyperbolic tangent function

```julia
T = Float32
x = Var(rand(T,100,10))
f = LSTM(T, 100, 100)
h = f(x)
```
"""
function LSTM(::Type{T}, insize::Int, hsize::Int, nlayers::Int, droprate::Float64, bidirectional::Bool;
    init_W=Normal(0,0.001), init_U=Orthogonal(), init_b=Fill(0), init_h=Fill(0), init_c=Fill(0)) where T

    params = []
    coef = bidirectional ? 2 : 1
    for l = 1:nlayers
        for _ = 1:coef
            Ws = []
            Us = []
            for i = 1:4
                s = l == 1 ? insize : hsize*coef
                push!(Ws, init_W(T,s,hsize))
                push!(Us, init_U(T,hsize,hsize))
            end
            W = param(cat(Ws...,dims=2))
            U = param(cat(Us...,dims=2))
            b = param(init_b(T,4hsize))
            h = param(init_h(T,hsize))
            c = param(init_c(T,hsize))
            push!(params, (W,U,b,h,c))
        end
    end
    LSTM(insize, hsize, nlayers, droprate, bidirectional, params)
end

function (lstm::LSTM)(x::Var, batchdims::Vector{Int})
    configure!(getparams(lstm)...)
    if iscpu()
        h = x
        coef = lstm.bidirectional ? 2 : 1
        for l = 1:lstm.nlayers
            i = (l-1) * coef + 1
            p = lstm.params[i]
            h1 = lstm_tstep(h, batchdims, p..., false)
            if lstm.bidirectional
                p = lstm.params[i+1]
                h2 = lstm_tstep(h, batchdims, p..., true)
                h = concat(1, h1, h2)
            else
                h = h1
            end
        end
        h
    elseif iscuda()
        lstm_cudnn(lstm, xs)
    else
        throw("Invalid backend.")
    end
end
(lstm::LSTM)(x::Node, batchdims) = Node(lstm, x, batchdims)

function lstm_tstep(x::Var, batchsize::Vector{Int}, W::Var, U::Var, b::Var, h0::Var, c0::Var, rev::Bool)
    @assert sum(batchsize) == size(x,2)
    WU = concat(1, W, U)

    cumdims = Array{Int}(undef, length(batchsize)+1)
    cumdims[1] = 1
    for i = 1:length(batchsize)
        cumdims[i+1] = cumdims[i] + batchsize[i]
    end
    perm = sortperm(batchsize, rev=true)

    hsize = length(h0)
    ht = concat(2, [h0 for i=1:length(batchsize)]...)
    ct = concat(2, [c0 for i=1:length(batchsize)]...)
    hts = Array{Var}(undef, size(x,2))
    cts = Array{Var}(undef, size(x,2))
    for t = 1:batchsize[perm[1]]
        xts = Var[]
        for p in perm
            t > batchsize[p] && break
            i = cumdims[p]
            i += rev ? batchsize[p]-t : t-1
            push!(xts, x[:,i:i])
        end
        xt = concat(2, xts...)
        if size(ht,2) > size(xt,2)
            ht = ht[:,1:size(xt,2)]
            ct = ct[:,1:size(xt,2)]
        end
        ht, ct = lstm_onestep(xt, WU, b, ht, ct)
        for j = 1:length(perm)
            p = perm[j]
            t > batchsize[p] && break
            i = cumdims[p]
            i += rev ? batchsize[p]-t : t-1
            hts[i] = ht[:,j:j]
            cts[i] = ct[:,j:j]
        end
    end
    concat(2, hts...)
end

function lstm_onestep(xt::Var, WU::Var, b::Var, ht::Var, ct::Var)
    a = linear(concat(1,xt,ht), WU, b)
    n = size(a,1) ÷ 4
    i = sigmoid(a[1:n,:])
    f = sigmoid(a[n+1:2n,:])
    ct = dot(f,ct) + dot(i,tanh(a[2n+1:3n,:]))
    o = sigmoid(a[3n+1:4n,:])
    ht = dot(o, tanh(ct))
    ht, ct
end

function batchsort(x::UniArray, batchsize::Vector{Int})
    perm = sortperm(batchsize, rev=true)
    cumdims = cumsum(batchsize)
    front = Base.front(size(x))
    view(x)
end

function lstm_cudnn(lstm::LSTM, x::Var, batchsize::Var)
    Ws = Var[]
    h0s = Var[]
    c0s = Var[]
    for (W,U,b,h,c) in lstm.params
        push!(Ws, vec(W), vec(U))
    end
    for (W,U,b,h,c) in lstm.params
        push!(Ws, b, b)
        push!(h0s, h)
        push!(c0s, c)
    end
    W = concat(1, Ws...)
    h0 = concat(1, h0s...)
    c0 = concat(1, c0s...)

    xs = unsafe_split(x.data, batchsize)
    perm = sortperm(batchsize, rev=true)
    t_x, t_batchsize = transpose_batch(xs[perm])

    dir = lstm.bidirectional ? CUDNN.CUDNN_BIDIRECTIONAL : CUDNN.CUDNN_UNIDIRECTIONAL
    mode = CUDNN.CUDNN_LSTM
    t_y, work = CUDNN.rnn(lstm.insize, lstm.hsize, lstm.nlayers, lstm.droprate, dir, mode,
        W.data, t_x, t_batchsize, istrain())

    y, _ = transpose_batch(unsafe_split(t_y,t_batchsize))
    ys = unsafe_split(y, batchsize[perm])
    y = cat(ys[perm]...,dims=ndims(y))
    Var(y, (lstm,x,batchsize,work,W))
end

function addgrad!(y::Var, lstm::LSTM, x::Var, batchsize::Vector{Int}, work, w::Var)
    t_gy, t_batchsize = transpose_batch(y.grad, batchsize)
    t_gx = CUDNN.∇rnn_data(work, t_gy) # this call is required for ∇rnn_weights!
    gx, _ = transpose_batch(t_gx, t_batchsize)
    isvoid(x.grad) || addto!(x.grad, gx)
    isvoid(w.grad) || CUDNN.∇rnn_weights!(work, w.grad)
end

@generated function transpose_batch(xs::Vector{CuMatrix{T}}) where T
    Ct = cstring(T)
    k = Kernel("""
    __global__ void transpose_batch(int n, $Ct *y, int *cumdimsY, $Ct **xs, int *cumdimsX) {
        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx >= n*cumdimsY[1]*cumdimsX[1]) return;

        int vj = idx / n;
        int vi = idx - vj * n;
        int j = vj / cumdimsY[1];
        int i = vj - j * cumdimsY[1];
        if (cumdimsY[j] + i < cumdimsY[j+1]) {
            int idxY = (cumdimsY[j] + i) * n + vi;
            int idxX = j * n + vi;
            y[idxY] = xs[i][idxX];
        }
    }""")
    quote
        batchsize_x = Array{Int}(length(xs))
        for i = 1:length(xs)
            batchsize_x[i] = size(xs[i], 2)
        end
        batchsize_y = transpose_dims(batchsize_x)
        cumdims_x = CuArray(cumsum_cint(batchsize_x))
        cumdims_y = CuArray(cumsum_cint(batchsize_y))
        y = CuArray{T}(size(xs[1],1), sum(batchsize_x))
        p_xs = CuArray(map(pointer,xs))
        gdims, bdims = cudims(size(xs[1],1)*batchsize_y[1]*batchsize_x[1])
        $k(gdims, bdims, size(xs[1],1), pointer(y), pointer(cumdims_y), pointer(p_xs), pointer(cumdims_x))
        y, batchsize_y
    end
end

@generated function transpose_batch2(x::CuMatrix{T}, batchsize_x::Vector{Int}) where T
    Ct = cstring(T)
    k = Kernel("""
    __global__ void transpose_batch(int n, $Ct *y, int *cumdimsY, $Ct *x, int *cumdimsX) {
        int idx = blockIdx.x * blockDim.x + threadIdx.x;
        if (idx >= n*cumdimsY[1]*cumdimsX[1]) return;

        int vj = idx / n;
        int vi = idx - vj * n;
        int j = vj / cumdimsY[1];
        int i = vj - j * cumdimsY[1];
        if (cumdimsY[j] + i < cumdimsY[j+1]) {
            int idxY = (cumdimsY[j] + i) * n + vi;
            int idxX = (cumdimsX[i] + j) * n + vi;
            y[idxY] = x[idxX];
        }
    }""")
    quote
        batchsize_y = transpose_dims(batchsize_x)
        cumdims_x = CuArray(cumsum_cint(batchsize_x))
        cumdims_y = CuArray(cumsum_cint(batchsize_y))

        y = similar(x)
        gdims, bdims = cudims(size(x,1)*batchsize_y[1]*batchsize_x[1])
        $k(gdims, bdims, size(x,1), pointer(y), pointer(cumdims_y), pointer(x), pointer(cumdims_x))
        y, batchsize_y
    end
end


sdt = sdt_indices(6, 7, 8, 9, adjusted=false)
@test sdt["dprime"] ≈ -0.023 atol=0.001
@test sdt["aprime"] ≈ 0.491 atol=0.001
@test sdt["beta"] ≈ 0.996 atol=0.001
@test sdt["c"] ≈ 0.169 atol=0.001
@test sdt["c_relative"] ≈ -7.43 atol=0.001
@test sdt["Xc"] ≈ 0.157 atol=0.001
@test sdt["bpp"] ≈ 0.263 atol=0.001

@test sdt["dprime_glm"] ≈ -0.0227 atol=0.001
@test sdt["Xc_glm"] ≈ 0.157 atol=0.001

sdt = sdt_indices(6, 7, 8, 9, adjusted=true)
@test sdt["dprime"] ≈ -0.023 atol=0.001
@test sdt["beta"] ≈ 0.995 atol=0.001

sdt = sdt_indices.([9, 1], [1, 9], [1, 9], [9, 1], adjusted=false)
@test length(sdt) == 2
@test sdt[1]["aprime"] ≈ 0.944 atol=0.001
@test sdt[1]["pr"] ≈ 0.800 atol=0.001
@test sdt[1]["br"] ≈ 0.500 atol=0.001
@test sdt[1]["hit_rate"] ≈ 0.900 atol=0.001
@test sdt[1]["fa_rate"] ≈ 0.100 atol=0.001


module BoxSymmetries
export BoxSym
export unit, inverse, ∘

using Random: AbstractRNG, Random
using ArgCheck
################################################################################
#### Perm
################################################################################
"""
Permutation
"""
struct Perm{N}
    perm::NTuple{N,Int}
    function Perm(perm::NTuple{N,Int}) where {N}
        @argcheck ispermutation(perm)
        return new{N}(perm)
    end
end
function Perm(args::Integer...)
    Perm(Int.(args))
end
function Base.show(io::IO, p::Perm)
    print(io, Perm, p.perm)
end
function ispermutation(perm::NTuple{N,Int}) where {N}
    ret = true
    for key in ntuple(identity, Val(N))
        ret = ret & (key in perm)
    end
    ret
end
ispermutation(perm::Perm) = true
function act_tuple(p::Perm{N}, t::NTuple{N}) where {N}
    map(p.perm) do i
        @inbounds t[i]
    end
end
function unit(::Type{Perm{N}}) where {N}
    Perm(ntuple(identity, Val(N)))
end
function Base.:(∘)(p1::Perm{N}, p2::Perm{N})::Perm{N} where {N}
    Perm(act_tuple(p1, p2.perm))
end
function inverse(p::Perm{N})::Perm{N} where {N}
    Perm(map(ntuple(identity, Val(N))) do i
        indexof(p.perm, i)
    end)
end
function indexof(t::NTuple, item)
    for i in eachindex(t)
        if t[i] === item
            return i
        end
    end
    error()
end
const PERM1 = [Perm(1)]
const PERM2 = [Perm(1,2), Perm(2,1)]
const PERM3 = [Perm(1,2,3),
               Perm(1,3,2), Perm(3,2,1), Perm(2,1,3),
               Perm(2,3,1), Perm(3,1,2),
]
Base.instances(::Type{Perm{1}}) = PERM1
Base.instances(::Type{Perm{2}}) = PERM2
Base.instances(::Type{Perm{3}}) = PERM3

################################################################################
#### BoxSym
################################################################################
struct BoxSym{N}
    # BoxSym is the semidirect product Boolⁿ ⋊ Sₙ
    # here axesperm ∈ Sₙ controls the induced action on the set of coordinate axes
    # flipsign controls for each axis, whether the sign is flipped. E.g. whether
    # eᵢ ↦ eⱼ or eᵢ ↦ -eⱼ for the standard basis
    axesperm::Perm
    flipsign::NTuple{N,Bool}
end
Base.ndims(o::BoxSym{N}) where {N} = N
Base.ndims(::Type{BoxSym{N}}) where {N} = N

function outaxes(o::BoxSym, axes::Tuple)
    @argcheck length(axes) == ndims(o)
    act_tuple(o.axesperm, axes)
end

isunit(g) = (g === unit(typeof(g)))

act_array!(out, g::BoxSym, x) = act_array_generic!(out, g, x)
act_array(g::BoxSym, x) = act_array_generic(g, x)

function act_array_generic!(out::AbstractArray{<:Any, N}, o::BoxSym{N}, x) where {N}
    @argcheck axes(out) == outaxes(o, axes(x))
    flipped_axes = map(flipaxis, axes(out), o.flipsign)
    @inbounds for I in CartesianIndices(x)
        inds = Tuple(I)
        inds = act_tuple(o.axesperm, inds)
        inds = map(getindex,flipped_axes, inds)
        Iout = CartesianIndex(inds)
        out[Iout] = x[I]
    end
    out
end
function act_array_generic(o::BoxSym, x)
    out = similar(x, outaxes(o,axes(x)))
    act_array_generic!(out, o, x)
end

function (o::BoxSym)(x)
    act_array(o,x)
end

function flipaxis(ax::AbstractRange, flip::Bool)
    if flip
        last(ax):-step(ax):first(ax)
    else
        first(ax):step(ax):last(ax)
    end
end

"""

    BoxSym(args::Integer...)

Create a box symmetry, according to args.
"""
function BoxSym(args::Integer...)
    perm = Perm(map(Int∘abs, args))
    flip = map(<(0), args)
    ret = BoxSym(perm, flip)
    ret
end
function astuple(o::BoxSym)
    (-1) .^ (o.flipsign) .* o.axesperm.perm
end
function Base.show(io::IO, o::BoxSym)
    print(io, "BoxSym", astuple(o))
end

function unit(::Type{BoxSym{N}})::BoxSym{N} where {N}
    BoxSym(ntuple(identity, Val(N))...)
end
function inverse(g::BoxSym{N})::BoxSym{N} where {N}
    n = g.flipsign
    h = g.axesperm
    h⁻¹ = inverse(h)
    n⁻¹ = n
    BoxSym{N}(
        h⁻¹,
        act_tuple(h⁻¹, n⁻¹),
   )
end
function Base.:(∘)(g1::BoxSym{N}, g2::BoxSym{N})::BoxSym{N} where {N}
    # We have
    # BoxSym = Boolⁿ ⋊ Sₙ
    # and we use the semidirect product formular
    n1 = g1.flipsign
    n2 = g2.flipsign
    h1 = g1.axesperm
    h2 = g2.axesperm
    h = h1∘h2
    n = map(⊻, n1, act_tuple(h1, n2))
    BoxSym{N}(h,n)
end

################################################################################
#### ALIASES
################################################################################
const BOXSYM_FROM_ALIAS2D = Dict(
    :unit      => BoxSym( 1, 2),
    :rot90     => BoxSym(-2, 1),
    :rot180    => BoxSym(-1,-2),
    :rot270    => BoxSym( 2,-1),
    :flipx     => BoxSym(-1, 2),
    :flipy     => BoxSym( 1,-2),
    :flipdiag  => BoxSym( 2, 1),
    :flipadiag => BoxSym(-2,-1),
)

BoxSym{2}(alias::Symbol) = BOXSYM_FROM_ALIAS2D[alias]

# TODO 3d aliases

################################################################################
#### instances
################################################################################

const SYMMETRIES1D = [BoxSym(1), BoxSym(-1)]
const SYMMETRIES2D = [
    BoxSym(1,2), BoxSym(1,-2), BoxSym(-1,2), BoxSym(-1,-2),
    BoxSym(2,1), BoxSym(2,-1), BoxSym(-2,1), BoxSym(-2,-1),
]
const SYMMETRIES3D = [
    BoxSym(1,2,3), BoxSym(1,2,-3), BoxSym(1,-2,3), BoxSym(1,-2,-3), BoxSym(-1,2,3), BoxSym(-1,2,-3), BoxSym(-1,-2,3), BoxSym(-1,-2,-3),
    BoxSym(1,3,2), BoxSym(1,3,-2), BoxSym(1,-3,2), BoxSym(1,-3,-2), BoxSym(-1,3,2), BoxSym(-1,3,-2), BoxSym(-1,-3,2), BoxSym(-1,-3,-2),
    BoxSym(3,2,1), BoxSym(3,2,-1), BoxSym(3,-2,1), BoxSym(3,-2,-1), BoxSym(-3,2,1), BoxSym(-3,2,-1), BoxSym(-3,-2,1), BoxSym(-3,-2,-1),
    BoxSym(2,1,3), BoxSym(2,1,-3), BoxSym(2,-1,3), BoxSym(2,-1,-3), BoxSym(-2,1,3), BoxSym(-2,1,-3), BoxSym(-2,-1,3), BoxSym(-2,-1,-3),
    BoxSym(2,3,1), BoxSym(2,3,-1), BoxSym(2,-3,1), BoxSym(2,-3,-1), BoxSym(-2,3,1), BoxSym(-2,3,-1), BoxSym(-2,-3,1), BoxSym(-2,-3,-1),
    BoxSym(3,1,2), BoxSym(3,1,-2), BoxSym(3,-1,2), BoxSym(3,-1,-2), BoxSym(-3,1,2), BoxSym(-3,1,-2), BoxSym(-3,-1,2), BoxSym(-3,-1,-2),
]

Base.instances(::Type{BoxSym{1}}) = SYMMETRIES1D
Base.instances(::Type{BoxSym{2}}) = SYMMETRIES2D
Base.instances(::Type{BoxSym{3}}) = SYMMETRIES3D

Random.rand(rng::AbstractRNG, G::Type{<:BoxSym}) = rand(rng, instances(G))

################################################################################
#### Precompile
################################################################################
let
    BoxSym(1)(fill(1.0, 1))
    BoxSym(1,-2)(fill(1.0, 1,1))
    BoxSym(1,-2,3)(fill(1.0, 1,1,1))
end

# TODO lazy iterator of all symmetries of given dimension
# TODO isrotation for checking if orientation is preserved
end #module


# Test linearity of sources
# Author: Philipp Witte, pwitte@eos.ubc.ca
# Date: January 2017
#

using PyCall, PyPlot, opesciSLIM.TimeModeling

## Set up model structure
n = (120,100)	# (x,y,z) or (x,z)
d = (10.,10.)
o = (0.,0.)

# Velocity [km/s]
v = ones(Float32,n) + 0.4f0
v[:,Int(round(end/2)):end] = 4f0

# Slowness squared [s^2/km^2]
m = (1f0./v).^2

# Setup info and model structure
nsrc = 1
ntComp = 250
info = Info(prod(n), nsrc, ntComp)	# number of gridpoints, number of experiments, number of computational time steps
model = Model(n,d,o,m)

## Set up receiver geometry
nxrec = 120
xrec = linspace(50f0,1150f0,nxrec)
yrec = 0f0
zrec = linspace(50f0,50f0,nxrec)

# receiver sampling and recording time
timeR = 1000f0	# receiver recording time [ms]
dtR = 4f0	# receiver sampling interval

# Set up receiver structure
recGeometry = Geometry(xrec,yrec,zrec;dt=dtR,t=timeR,nsrc=nsrc)

## Set up source geometries (cell array with source locations for each shot)
xsrc1 = 300f0
ysrc1 = 0f0
zsrc1 = 20f0

xsrc2 = 600f0
ysrc2 = 0f0
zsrc2 = 20f0

# source sampling and number of time steps
timeS = 1000f0
dtS = 4f0	# receiver sampling interval

# Set up source structure
srcGeometry1 = Geometry(xsrc1,ysrc1,zsrc1;dt=dtS,t=timeS)
srcGeometry2 = Geometry(xsrc2,ysrc2,zsrc2;dt=dtS,t=timeS)

# setup wavelet
f0 = 0.01f0
wavelet = source(timeS,dtS,f0)

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

# Modeling operators
Pr = joProjection(info,recGeometry)
Ps1 = joProjection(info,srcGeometry1)
Ps2 = joProjection(info,srcGeometry2)
F = joModeling(info,model)
q1 = joData(srcGeometry1,wavelet)
q2 = joData(srcGeometry2,wavelet)

d1 = Pr*F*Ps1'*q1
d2 = Pr*F*Ps2'*q2
d3 = Pr*F*(Ps1'*q1 + Ps2'*q2)
d4 = Pr*F*(Ps1'*q1 - Ps2'*q2)

println("addition: ", norm(d3 - (d1 + d2)))
println("subtraction: ", norm(d4 - (d1 - d2)))


using Test
using PolynomialBasis
using ImplicitDomainQuadrature
# using Revise
using CutCellDG
include("../useful_routines.jl")

polyorder = 1
numqp = 2

elasticitybasis = LagrangeTensorProductBasis(2, polyorder)
levelsetbasis = HermiteTensorProductBasis(2)

x0 = [0.0, 0.0]
meshwidths = [3.0, 1.0]
nelements = [3, 1]
points = interpolation_points(elasticitybasis)
dim, nf = size(interpolation_points(levelsetbasis))
dgmesh = CutCellDG.DGMesh(x0, meshwidths, nelements, points)
cgmesh = CutCellDG.CGMesh(x0, meshwidths, nelements, nf)

normal = [1.0, 0.0]
xI = [0.5, 0.0]
levelset = CutCellDG.LevelSet(
    x -> plane_distance_function(x, normal, xI),
    cgmesh,
    levelsetbasis,
)
cutmesh = CutCellDG.CutMesh(dgmesh, levelset)

interfacequads = CutCellDG.InterfaceQuadratures(cutmesh, levelset, numqp)
@test size(interfacequads.quads) == (2, 1)
testcelltoquad = [1, 0, 0]
@test allequal(interfacequads.celltoquad, testcelltoquad)

testnormals = repeat(normal, inner = (1, numqp))
@test allapprox(
    testnormals,
    CutCellDG.interface_normals(interfacequads, 1),
    1e3eps(),
)
@test interfacequads[1, 1] ≈ interfacequads[-1, 1]

testtangents = repeat([0.0, 1.0], inner = (1, numqp))
testscaleareas = repeat([0.5], numqp)
@test allapprox(
    testtangents,
    CutCellDG.interface_tangents(interfacequads, 1),
    1e3eps(),
)
@test allapprox(
    testscaleareas,
    CutCellDG.interface_scale_areas(interfacequads, 1),
)


##Test ABC algorithms and related functions.
##This is done by running a simple normal example.
##It is simply tested that these do not produce errors.
using ABCDistances
using Distributions

##Set up abcinput
function sample_sumstats(pars::Array)
    success = true
    stats = [pars[1] + 0.1*randn(1); randn(1)]
    (success, stats)
end

sobs = [0.0,0.0]

abcinput = ABCInput();
abcinput.prior = MvNormal(1, 100.0);
abcinput.sample_sumstats = sample_sumstats;
abcinput.abcdist = WeightedEuclidean(sobs);
abcinput.nsumstats = 2;

##ABC rejection
srand(1)
out1 = abcRejection(abcinput, 1000, 200)
abcRejection(abcinput, 1000, 0.1)

##ABC-PMC (doesn't quite test all options)
srand(1)
out2 = abcPMC3V(abcinput, 200, 1/2, 5000, silent=true)
abcPMC5(abcinput, 200, 1/2, 5000, silent=true)
abcPMC3V(abcinput, 200, 1/2, 5000, store_init=true, silent=true)
abcPMC3V(abcinput, 200, 1/2, 5000, store_init=true, silent=true)

abcPMC3(abcinput, 200, 1/2, 5000, silent=true)
abcPMC3(abcinput, 200, 1/2, 5000, store_init=true, silent=true)
abcPMC3(abcinput, 200, 1/2, 5000, diag_perturb=true, silent=true)

abcPMC4(abcinput, 200, 1/2, 5000, silent=true);
abcPMC4(abcinput, 200, 1/2, 5000, store_init=true, silent=true);

##Mean and variance functions
parameter_means(out1)
parameter_means(out2)
parameter_vars(out1)
parameter_vars(out2)


using Catlab
using Catlab.WiringDiagrams
using Catlab.Doctrines
import Catlab.Doctrines: ⊗, id
import Base: ∘
include("SIR-simple__functions.jl")
⊗(a::WiringDiagram, b::WiringDiagram) = otimes(a, b)
∘(a::WiringDiagram, b::WiringDiagram) = compose(b, a)
⊚(a,b) = b ∘ a

SIR_simple__global__sir__0__beta__neg1, SIR_simple__global__sir__0__S__neg1, SIR_simple__global__sir__0__I__neg1, SIR_simple__global__sir__0__R__neg1, SIR_simple__global__sir__0__dt__neg1, SIR_simple__global__sir__0__gamma__neg1 = Ob(FreeSymmetricMonoidalCategory, :SIR_simple__global__sir__0__beta__neg1, :SIR_simple__global__sir__0__S__neg1, :SIR_simple__global__sir__0__I__neg1, :SIR_simple__global__sir__0__R__neg1, :SIR_simple__global__sir__0__dt__neg1, :SIR_simple__global__sir__0__gamma__neg1)
SIR_simple__global__sir__0__infected__0, SIR_simple__global__sir__0__recovered__0 = Ob(FreeSymmetricMonoidalCategory, :SIR_simple__global__sir__0__infected__0, :SIR_simple__global__sir__0__recovered__0)
SIR_simple__global__sir__0__I__0, SIR_simple__global__sir__0__S__0, SIR_simple__global__sir__0__R__0 = Ob(FreeSymmetricMonoidalCategory, :SIR_simple__global__sir__0__I__0, :SIR_simple__global__sir__0__S__0, :SIR_simple__global__sir__0__R__0)

id_SIR_simple__global__sir__0__R__neg1 = id(Ports([SIR_simple__global__sir__0__R__neg1]))
id_SIR_simple__global__sir__0__S__neg1 = id(Ports([SIR_simple__global__sir__0__S__neg1]))
id_SIR_simple__global__sir__0__I__neg1 = id(Ports([SIR_simple__global__sir__0__I__neg1]))

IN_0 = WiringDiagram(Hom(:L0_REWIRE, SIR_simple__global__sir__0__beta__neg1 ⊗ SIR_simple__global__sir__0__S__neg1 ⊗ SIR_simple__global__sir__0__I__neg1 ⊗ SIR_simple__global__sir__0__R__neg1 ⊗ SIR_simple__global__sir__0__dt__neg1 ⊗ SIR_simple__global__sir__0__gamma__neg1, SIR_simple__global__sir__0__beta__neg1 ⊗ SIR_simple__global__sir__0__S__neg1 ⊗ SIR_simple__global__sir__0__I__neg1 ⊗ SIR_simple__global__sir__0__R__neg1 ⊗ SIR_simple__global__sir__0__dt__neg1 ⊗ SIR_simple__global__sir__0__gamma__neg1 ⊗ SIR_simple__global__sir__0__I__neg1 ⊗ SIR_simple__global__sir__0__dt__neg1 ⊗ SIR_simple__global__sir__0__R__neg1 ⊗ SIR_simple__global__sir__0__S__neg1 ⊗ SIR_simple__global__sir__0__I__neg1))
WD_SIR_simple__global__sir__0__recovered__0 = WiringDiagram(Hom(SIR_simple__sir__assign__recovered__0, SIR_simple__global__sir__0__gamma__neg1 ⊗ SIR_simple__global__sir__0__I__neg1 ⊗ SIR_simple__global__sir__0__dt__neg1, SIR_simple__global__sir__0__recovered__0))
WD_SIR_simple__global__sir__0__infected__0 = WiringDiagram(Hom(SIR_simple__sir__assign__infected__0, SIR_simple__global__sir__0__beta__neg1 ⊗ SIR_simple__global__sir__0__S__neg1 ⊗ SIR_simple__global__sir__0__I__neg1 ⊗ SIR_simple__global__sir__0__R__neg1 ⊗ SIR_simple__global__sir__0__dt__neg1, SIR_simple__global__sir__0__infected__0))
OUT_1 = IN_0 ⊚ (WD_SIR_simple__global__sir__0__infected__0 ⊗ WD_SIR_simple__global__sir__0__recovered__0 ⊗ id_SIR_simple__global__sir__0__R__neg1 ⊗ id_SIR_simple__global__sir__0__S__neg1 ⊗ id_SIR_simple__global__sir__0__I__neg1)

IN_1 = WiringDiagram(Hom(:L1_REWIRE, SIR_simple__global__sir__0__infected__0 ⊗ SIR_simple__global__sir__0__recovered__0 ⊗ SIR_simple__global__sir__0__R__neg1 ⊗ SIR_simple__global__sir__0__S__neg1 ⊗ SIR_simple__global__sir__0__I__neg1, SIR_simple__global__sir__0__I__neg1 ⊗ SIR_simple__global__sir__0__infected__0 ⊗ SIR_simple__global__sir__0__recovered__0 ⊗ SIR_simple__global__sir__0__S__neg1 ⊗ SIR_simple__global__sir__0__infected__0 ⊗ SIR_simple__global__sir__0__R__neg1 ⊗ SIR_simple__global__sir__0__recovered__0))
WD_SIR_simple__global__sir__0__R__0 = WiringDiagram(Hom(SIR_simple__sir__assign__R__0, SIR_simple__global__sir__0__R__neg1 ⊗ SIR_simple__global__sir__0__recovered__0, SIR_simple__global__sir__0__R__0))
WD_SIR_simple__global__sir__0__S__0 = WiringDiagram(Hom(SIR_simple__sir__assign__S__0, SIR_simple__global__sir__0__S__neg1 ⊗ SIR_simple__global__sir__0__infected__0, SIR_simple__global__sir__0__S__0))
WD_SIR_simple__global__sir__0__I__0 = WiringDiagram(Hom(SIR_simple__sir__assign__I__0, SIR_simple__global__sir__0__I__neg1 ⊗ SIR_simple__global__sir__0__infected__0 ⊗ SIR_simple__global__sir__0__recovered__0, SIR_simple__global__sir__0__I__0))
OUT_2 = OUT_1 ⊚ IN_1 ⊚ (WD_SIR_simple__global__sir__0__I__0 ⊗ WD_SIR_simple__global__sir__0__S__0 ⊗ WD_SIR_simple__global__sir__0__R__0)


#=
Title: Discrete Simulation of 2-Dimensional Light Paths in Continuous Media
Author:
=#

#finds the magnitude of a vector in 2-space
function magnitude(vector) #vector type list
    return sqrt(vector[1]^2+vector[2]^2)
end

#finds the angle between two vectors
function get_angle(a,b)#a,b are type list
    if magnitude(a)>magnitude(b)
        hypotenuse = magnitude(a)
        adjacent = magnitude(b)
    else
        hypotenuse = magnitude(b)
        adjacent = magnitude(a)
    end
    if hypotenuse == 0
        hypotenuse = 0.00001
        break
    end
    return acos(adjacent / hypotenuse)
end

#finds the lowest angle between the norm and the velocity
function choose_angle(normal,velocity)
    positive = get_angle(normal,velocity)
    negative = get_angle(-normal,velocity)
    if abs(positive)<=abs(negative)
        return positive
    else
        return negative
    end
end

#rotates a matrix by some angle in radians
function instantaneous_rotation(current_velocity,current_angle)
    x = velocity[1]*cos(angle)-velocity[2]*sin(angle)
    y = velocity[1]*sin(angle)+velocity[2]*cos(angle)
    return [x,y]
end

#the continuous refractive index at any point
function refractive_index(point) #point type list (x,y)
    return (point[1]^2)*(point[2]^2)+1
end

#Gradient of level plane W(x,y,n) is normal to surface n
#the particular normal at some point
function projected_normal(point) #point type list (x,y)
    x = 2*point[1]*point[2]^2
    y = 2*point[2]*point[1]^2
    return [x,y] #we only need the projection onto 2-space
end

#generates a random initial velocity
function generate_velocity()
    candidate = [rand(-c:c)-rand(),0]
    candidate[2] = sqrt(c^2 - candidate[1]^2)
    return candidate
end

#generates a random initial position
function generate_position()
    return [rand(-10:10)+rand(),rand(-10:10)+rand()]
end

#=
Given n1 = c/v1 and n2 = c/v2, n1/n2 = (c/v1)(v2/c) = (v2/v1)(c/c)= (v2/v1)
so (v1 n1)/n2 = v2. Or in this case, n1 is the last n(x,z), v1 is the current
velocity, and n2 is the current n(x,z), and v2 is the velocity we are trying
to calculate.
=#

#Main Body of Code

#initial conditions of the photon:
c = 2.99e8
current_position,current_velocity = generate_position(),generate_velocity()
current_n = c / magnitude(current_velocity)

running = true
while running
    #have to figure out how to update current_n
    incident_normal = projected_normal(current_position)
    incident_angle = choose_angle(incident_normal,current_velocity)
    #the refractive angle between the refractive velocity and the refractive normal
    refractive_angle = asin((current_n*sin(incident_angle))/refractive_index(current_position))
    #=
    given you want to start with some current velocity, this function gives
    you the factor by which to multiply the current velocity, in order to get the
    desired refractive velocity
    =#
    scaling_factor = c/(magnitude(current_velocity)*refractive_index(current_position))
    next_velocity = scaling_factor*current_velocity#temporary... not the real next_velocity. still needs to be rotated

    current_velocity,current_position = next_velocity,next_position
    current_n = refractive_index(current_position)
end

#=
If p in R^2 is a function of x and y, then we can parametrize
x and y so that x = x(t) and y = y(t). Then, if v is current_velocity
in R^2, and v is parametrized with t as well, so that v[1] = x'{t}
and v[2] = y'(t). Then x'(t) is defined as the limit as h->0 of
(x(t+h)-x(t))/h. If instead of 0, we allow h->epsilon, then
x'(t) = (x(t+epsilon)-x(t))/epsilon. We can similarly say for y'(t)
that y'(t) = (y(t+epsilon)-y(t))/epsilon.
Then epsilon*x'(t) = x(t+epsilon) - x(t) and
epsilon*y'(t) = y(t+epsilon) - y(t).
Then epsilon*x'(t) + x(t) = x(t+epsilon)
and epsilon*y'(t) + y(t) = y(t+epsilon).
=#

#=
General sketch:
*photon = [current_velocity,current_position]
*use current_velocity and current_position to find next_velocity.
*use next_velocity to find next_position
*assign: current_velocity,current_position = next_velocity,next_position

currents = [current_position,current_velocity,current_n,current_angle]
nexts = [next_position,next_velocity,refractive_index(current_position),
next_angle]

refractive_index(current_position) = c / magnitude(next_velocity)
magnitude(next_velocity) = c / refractive_index(current_position)

next_velocity is proportional to current_velocity. Therefore
magnitude(next_velocity) = k*magnitude(current_velocity)
=#


using LinearAlgebra: diagind

"""
    OuyangModelOne(A=0.01, B=1.6, Γ=1e-4, N=10, ΔE=1.6e-2, D=1.0)

Model #1 from [Ouyang and Subotnik](http://aip.scitation.org/doi/10.1063/1.4908032).
See also [Ouyang's thesis](https://repository.upenn.edu/edissertations/2508/).
"""
Parameters.@with_kw struct OuyangModelOne <: DiabaticFrictionModel
    A::Float64 = 0.01
    B::Float64 = 1.6
    Γ::Float64 = 1e-4
    N::Int = 10
    ΔE::Float64 = 1.6e-2
    ρ::Float64 = N / ΔE
    C::Float64 = sqrt(Γ / 2π*ρ) / N
    D::Float64 = 1.0

    n_states::Int = N + 1
end

NQCModels.nstates(model::OuyangModelOne) = model.n_states
NQCModels.ndofs(::OuyangModelOne) = 1

function NQCModels.potential!(model::OuyangModelOne, V::Hermitian, r::Real)
    Parameters.@unpack A, B, C, D, ΔE, N = model

    Vsys(x) = A*tanh(B*x)
    Vbath(x, Δ) = -Vsys(x) + Δ
    Vsb(x) = C*exp(-D*x^2)

    V[1,1] = Vsys(r)
    Δs = range(-ΔE/2, ΔE/2, length=N)
    V.data[1,2:N+1] .= Vsb(r)
    V[diagind(N+1,N+1)[2:end]] .= Vbath.(r, Δs)

    return V
end

function NQCModels.derivative!(model::OuyangModelOne, derivative::Hermitian, r::Real)
    Parameters.@unpack A, B, C, D, N = model

    Dsys(x) = B*A*sech(B*x)^2
    Dbath(x) = -Dsys(x)
    Dsb(x) = -2D*x*C*exp(-D*x^2)

    derivative[1,1] = Dsys(r)
    derivative.data[1,2:N+1] .= Dsb(r)
    derivative[diagind(N+1,N+1)[2:end]] .= Dbath(r)

    return derivative
end


using Parameters
using ModelParameters

Base.@kwdef mutable struct ExpuhNode{P, A<:Real} <: IHACRESNode
    @network_node

    # https://wiki.ewater.org.au/display/SD41/IHACRES-CMD+-+SRG
    d::P = Param(200.0, bounds=(10.0, 550.0))  # flow threshold
    d2::P = Param(2.0, bounds=(0.0001, 10.0))   # flow threshold, multiplier applied to d
    e::P = Param(1.0, bounds=(0.1, 1.5))  # temperature to PET conversion factor
    f::P = Param(0.8, bounds=(0.01, 3.0))  # plant stress threshold factor (multiplicative factor of d)
    tau_q::P = Param(1.0, bounds=(0.0, 5.0))
    tau_s::P = Param(5, bounds=(5.0, 200.0))
    v_s::P = Param(0.5, bounds=(0.0, 1.0))

    storage_coef::P = Param(2.9, bounds=(0.2, 10.0))

    level_params::Array{P, 1} = [
        Param(-0.01, bounds=(-10.0, -0.01)),  # p1
        Param(0.8, bounds=(0.0, 1.5)),  # p2
        Param(4.5, bounds=(0.0, 20.0)), # p3
        Param(5.0, bounds=(1.0, 10.0)), # p4
        Param(0.35, bounds=(0.0, 1.0)), # p5
        Param(1.41, bounds=(-2.0, 2.0)), # p6
        Param(-1.45, bounds=(-2.5, 0.0)), # p7
        Param(6.75, bounds=(0.0, 10.0)), # p8
        Param(150.0, bounds=(50.0, 200.0)) # ctf
    ]

    storage::Array{A} = [100.0]
    quick_store::Array{A} = [0.0]
    slow_store::Array{A} = [0.0]
    outflow::Array{A} = []
    effective_rainfall::Array{A} = []
    et::Array{A} = []
    inflow::Array{A} = []
    level::Array{A} = []
    gw_store::Array{A} = [0.0]
end

function ExpuhNode(name::String, spec::Dict)
    n = create_node(ExpuhNode, name, spec["area"])

    node_params = spec["parameters"]
    n_lparams = n.level_params
    s_lparams = spec["level_params"]
    node_params["level_params"] = Param[
        Param(s_lparams[1], bounds=n_lparams[1].bounds)
        Param(s_lparams[2], bounds=n_lparams[2].bounds)
        Param(s_lparams[3], bounds=n_lparams[3].bounds)
        Param(s_lparams[4], bounds=n_lparams[4].bounds)
        Param(s_lparams[5], bounds=n_lparams[5].bounds)
        Param(s_lparams[6], bounds=n_lparams[6].bounds)
        Param(s_lparams[7], bounds=n_lparams[7].bounds)
        Param(s_lparams[8], bounds=n_lparams[8].bounds)
        Param(s_lparams[9], bounds=n_lparams[9].bounds)
    ]

    node_params["storage"] = [node_params["initial_storage"]]
    delete!(node_params, "initial_storage")

    for (k, p) in node_params
        s = Symbol(k)
        if p isa String
            p = eval(Meta.parse(p))
        end

        try
            f = getfield(n, s)
            setfield!(n, s, Param(p, bounds=f.bounds))
        catch err
            msg = sprint(showerror, err, catch_backtrace())
            if occursin("no field bounds", msg)
                setfield!(n, s, p)
            else
                throw(err)
            end
        end
    end

    return n
end

function ExpuhNode(name::String, area::Float64, d::Float64, d2::Float64, e::Float64, f::Float64,
                    tau_q::Float64, tau_s::Float64, v_s::Float64, s_coef::Float64,
                    store::Float64, quick::Float64, slow::Float64)
    return ExpuhNode{Param, Float64}(
        name=name,
        area=area,
        d=d,
        d2=d2,
        e=e,
        f=f,
        tau_q=tau_q,
        tau_s=tau_s,
        v_s=v_s,
        storage_coef=s_coef,
        storage=[store],
        quick_store=[quick],
        slow_store=[slow]
    )
end

"""
    run_node!(s_node::ExpuhNode, rain::Float64, evap::Float64,
              inflow::Float64, ext::Float64, gw_exchange::Float64;
              current_store=nothing,
              quick_store=nothing,
              slow_store=nothing,
            )::Tuple

Run given IHACRES ExpuhNode for a time step based on last known state.
"""
function run_node!(s_node::ExpuhNode,
                   rain::Float64,
                   evap::Float64,
                   inflow::Float64,
                   ext::Float64,
                   gw_exchange::Float64;
                   current_store=nothing,
                   quick_store=nothing,
                   slow_store=nothing,
                )::Tuple

    if !isnothing(current_store)
        s_node.storage[end] = current_store
    end

    if !isnothing(quick_store)
        s_node.quick_store[end] = quick_store
    end

    if !isnothing(slow_store)
        s_node.slow_store[end] = slow_store
    end

    cmd::Float64 = s_node.storage[end]

    e_rainfall = @ccall IHACRES.calc_effective_rainfall(rain::Cdouble, cmd::Cdouble, s_node.d::Cdouble, s_node.d2::Cdouble)::Cdouble
    mf = @ccall IHACRES.calc_trig_interim_cmd(cmd::Cdouble, s_node.d::Cdouble, e_rainfall::Cdouble)::Cdouble

    et::Float64 = @ccall IHACRES.calc_ET(s_node.e::Cdouble, evap::Cdouble, mf::Cdouble, s_node.f::Cdouble, s_node.d::Cdouble)::Cdouble
    cmd = @ccall IHACRES.calc_cmd(mf::Cdouble, rain::Cdouble, et::Cdouble, e_rainfall::Cdouble)::Cdouble

    (prev_q, prev_s) = (s_node.quick_store[end], s_node.slow_store[end])

    flow_res = [0.0, 0.0, 0.0]
    @ccall IHACRES.calc_flows(flow_res::Ptr{Cdouble}, prev_q::Cdouble, prev_s::Cdouble, s_node.v_s::Cdouble, e_rainfall::Cdouble,
                              s_node.area::Cdouble, s_node.tau_q::Cdouble, s_node.tau_s::Cdouble)::Cvoid
    (quick_store, slow_store, outflow) = flow_res

    gw_store = s_node.gw_store[end]
    routing_res = [0.0, 0.0]
    @ccall IHACRES.routing(
            routing_res::Ptr{Cdouble},
            gw_store::Cdouble,
            s_node.storage_coef::Cdouble,
            inflow::Cdouble,
            outflow::Cdouble,
            ext::Cdouble,
            gw_exchange::Cdouble)::Cvoid
    (gw_store, outflow) = routing_res

    level::Float64 = @ccall IHACRES.calc_ft_level(outflow::Cdouble, s_node.level_params::Ptr{Cdouble})::Cdouble

    update_state(s_node, cmd, e_rainfall, et, quick_store, slow_store, outflow, level, gw_store)

    return (outflow, level)
end

function run_node_with_temp!(s_node::ExpuhNode, rain::Float64, temp::Float64, inflow::Float64, ext::Float64; gw_exchange::Float64=0.0)::Tuple

    cmd::Float64 = s_node.storage[end]

    # rain = rain * s_node.area
    e_rainfall = @ccall IHACRES.calc_effective_rainfall(rain::Cdouble, cmd::Cdouble, s_node.d::Cdouble, s_node.d2::Cdouble)::Cdouble
    # mf = @ccall IHACRES.calc_linear_interim_cmd(cmd::Cdouble, s_node.d::Cdouble, e_rainfall::Cdouble)::Cdouble
    mf = @ccall IHACRES.calc_trig_interim_cmd(cmd::Cdouble, s_node.d::Cdouble, e_rainfall::Cdouble)::Cdouble

    et::Float64 = @ccall IHACRES.calc_ET_from_T(
        s_node.e::Cdouble,
        temp::Cdouble,
        mf::Cdouble,
        s_node.f::Cdouble,
        s_node.d::Cdouble
    )::Cdouble

    cmd = @ccall IHACRES.calc_cmd(current_store::Cdouble, rain::Cdouble, et::Cdouble, e_rainfall::Cdouble)::Cdouble

    (prev_q, prev_s) = (s_node.quick_store[end], s_node.slow_store[end])

    flow_res = [0.0, 0.0, 0.0]
    @ccall IHACRES.calc_flows(flow_res::Ptr{Cdouble}, prev_q::Cdouble, prev_s::Cdouble, s_node.v_s::Cdouble, e_rainfall::Cdouble,
                              s_node.area::Cdouble, s_node.tau_q::Cdouble, s_node.tau_s::Cdouble)::Cvoid
    (quick_store, slow_store, outflow) = flow_res

    gw_store = s_node.gw_store[end]
    routing_res = [0.0, 0.0]
    @ccall IHACRES.routing(
            routing_res::Ptr{Cdouble},
            gw_store::Cdouble,
            s_node.storage_coef::Cdouble,
            inflow::Cdouble,
            outflow::Cdouble,
            ext::Cdouble,
            gw_exchange::Cdouble)::Cvoid
    (gw_store, outflow) = routing_res

    level::Float64 = @ccall IHACRES.calc_ft_level(outflow::Cdouble, s_node.level_params::Ptr{Cdouble})::Cdouble

    update_state(s_node, cmd, e_rainfall, et, quick_store, slow_store, outflow, level, gw_store)

    return (outflow, level)
end

function update_params!(node::ExpuhNode, d::Float64, d2::Float64, e::Float64, f::Float64,
                        tau_q::Float64, tau_s::Float64, v_s::Float64, s_coef::Float64)::Nothing
    node.d = Param(d, bounds=node.d.bounds)
    node.d2 = Param(d2, bounds=node.d2.bounds)
    node.e = Param(e, bounds=node.e.bounds)
    node.f = Param(f, bounds=node.f.bounds)
    node.tau_q = Param(tau_q, bounds=node.tau_q.bounds)
    node.tau_s = Param(tau_s, bounds=node.tau_s.bounds)
    node.v_s = Param(v_s, bounds=node.v_s.bounds)
    node.storage_coef = Param(s_coef, bounds=node.storage_coef.bounds)

    return nothing
end


# f(x)=diagA*X
function FdiagAmultX(A::Array,X::Array)
    if size(A)==(1,1)
        return (A[1].*X,nothing)
    elseif size(X)==(1,1)
        return (diagm(vec(A)).*X[1],nothing)
    else
        return (diagm(vec(A))*X,nothing)
    end
end

function FdiagAmultX_inplace(value,auxvalue,A::Array,X::Array)
    if size(A)==(1,1)
        copy!(value,A[1]*X)
    elseif size(X)==(1,1)
        copy!(value,diagm(vec(A))*X[1])
    else
        copy!(value,diagm(vec(A))*X) # BLAS has no dgmm routine
    end
end

function DdiagAmultX(derivativeIDX,f_c,faux_c,grad_c,grad_n,A,X)
    if derivativeIDX==1
        if size(A)==(1,1)
            axpy!(1.0,[sum(X.*grad_c)],grad_n)
        elseif size(X)==(1,1)
            axpy!(X[1],sum(grad_c,2),grad_n)
        else
            axpy!(1.0,sum(X.*grad_c,2),grad_n)
        end
    elseif derivativeIDX==2
        if size(A)==(1,1)
            axpy!(A[1],grad_c,grad_n)
        elseif size(X)==(1,1)
            #axpy!(1.0,[sum(A.*diag(grad_c))],grad_n)
            axpy!(1.0,[sum(sum(A,2).*diag(grad_c))],grad_n)
        else
            axpy!(1.0,diagm(vec(A))*grad_c,grad_n)
        end
    end
end

if PROC=="GPU"
    function FdiagAmultX_inplace(value::CudaArray,auxvalue,A::CudaArray,X::CudaArray)
        if size(A)==(1,1)
            gax!(A,X,value)
        elseif size(X)==(1,1)
            diagm!(A,value)
            scale!(X,value) # nb argument converse of Base.scale!
        else
            CUBLAS.dgmm!('L',X,vec(A),value)
        end
    end

    function DdiagAmultX(derivativeIDX,f_c,faux_c,grad_c,grad_n,A::CudaArray,X::CudaArray)
        if derivativeIDX==1
            if size(A)==(1,1)
                tmp=CudaArray(Float64,size(X))
                vmult!(1.0,X,grad_c,tmp)
                sum_update!(1.0,tmp,1.0,grad_n)
                free(tmp)
            elseif size(X)==(1,1)
                tmp=CudaArray(Float64,size(grad_n))
                sum!(grad_c,2,0.0,tmp) # tmp=sum(grad_c,2)
                alphaaxpy!(1.0,X,tmp,grad_n)
                free(tmp)
            else
                tmp=CudaArray(Float64,size(X))
                vmult!(1.0,X,grad_c,tmp) #tmp=X.*grad_c
                sum!(tmp,2,1.0,grad_n)
            end
        elseif derivativeIDX==2
            if size(A)==(1,1)
                alphaaxpy!(1.0,A,grad_c,grad_n)
            elseif size(X)==(1,1)
              tmp1=CudaArray(Float64,(size(A,1),1))
              sum!(A,2,0.0,tmp1) # sum(A,2)
              tmp2=CudaArray(Float64,(size(A,1),1))
              diag!(grad_c,tmp2)
              vmult!(1.0,tmp1,tmp2,tmp1)
              sum_update!(1.0,tmp1,1.0,grad_n)
              free(tmp1); free(tmp2)
#            axpy!(1.0,[sum(sum(A,2).*diag(grad_c))],grad_n)
            else
                tmp=CudaArray(Float64,size(grad_c))
                             CUBLAS.dgmm!('L',grad_c,vec(A),tmp)
                axpy!(1.0,tmp,grad_n)
                free(tmp)
            end
        end
    end
end

Derivative[FdiagAmultX]=DdiagAmultX
export FdiagAmultX

Inplace[FdiagAmultX]=FdiagAmultX_inplace


using Test
using NiLang, NiLang.AD
using LinearAlgebra
using Random

@testset "inv" begin
    Random.seed!(2)

    id = [1 0 0; 0 1 0; 0 0 1.0]
    @test check_inv(i_inv!, (randn(3, 3), randn(3, 3)))
    @test check_inv(PlusEq(det), (0.3, randn(3, 3)))
    @test check_inv(PlusEq(logdet), (0.3, rand(3, 3) .+ id))
    @test check_grad(PlusEq(det), (0.3, randn(3, 3)), iloss=1)
    @test check_grad(PlusEq(logdet), (0.3, rand(3, 3) .+ id), iloss=1)

    @i function loss(out!, y, A)
        i_inv!(y, A)
        out! += y[1,1]
    end
    @test check_grad(loss, (0.0, randn(3, 3), randn(3, 3)); iloss=1)
end

@testset "affine" begin
    Random.seed!(2)
    A = randn(5, 5)
    b = randn(5)
    x = randn(5)
    y! = zeros(5)
    @test i_affine!(y!, A, b, x)[1] ≈ A*x + b
end


export isnorm, norm_equation

#TODO: verbose printing

function norm_1_generators(A::Vector{NfOrdIdl})
  @assert all(isprime, A)
  @assert all(x->x.gen_one == A[1].gen_one, A)

  f = matrix(FlintZZ, 1, length(A), [degree(x) for x = A])
  k = nullspace(f)[2]

  id = [FacElem(A, [k[i,j] for i=1:length(A)]) for j=1:ncols(k)]
  return id
end

@doc Markdown.doc"""
    norm_equation(K::AnticNumerField, a) -> nf_elem

For $a$ an integer or rational, try to find $T \in K$ s.th.
$N(T) = a$. Raises an error if unsuccessful.
"""
function norm_equation(K::AnticNumberField, a)
  fl, s = isnorm(K, a)
  if fl
    return evaluate(s)
  end
  error("no solution")
end

@doc Markdown.doc"""
    isnorm(K::AnticNumberField, a) -> Bool, nf_elem

For $a$ an integer or rational, try to find $T \in K$ s.th. $N(T) = a$
holds. If successful, return true and $T$, otherwise false and some element.
The element will be returned in factored form.
"""
function isnorm(K::AnticNumberField, a::Integer)
  return isnorm(K, fmpz(a))
end
function isnorm(K::AnticNumberField, a::fmpq)
  fl, s = isnorm(K, numerator(a)*denominator(a)^(degree(K)-1))
  return fl, s * FacElem(Dict(K(denominator(a)) => fmpz(-1)))
end
function isnorm(K::AnticNumberField, a::Rational)
  return isnorm(K, fmpq(a))
end

@doc Markdown.doc"""
    isnorm(K::AnticNumberField, a::fmpz; extra::Vector{fmpz}) -> Bool, nf_elem

For a fmpz $a$, try to find $T \in K$ s.th. $N(T) = a$
holds. If successful, return true and $T$, otherwise false and some element.
In \testtt{extra} one can pass in additional prime numbers that
are allowed to occur in the solution. This will then be supplemented.
The element will be returned in factored form.
"""
function isnorm(K::AnticNumberField, a::fmpz; extra::Vector{fmpz}=fmpz[])
  L = lll(maximal_order(K))
  C, mC = narrow_class_group(L)
#  println("narrow group is : $C")
  S = union(Set(keys(factor(a).fac)), Set(keys(factor(discriminant(L)).fac)))
  S = union(S, Set(extra))

  g = Set(elem_type(C)[])
  for p = S
    P = prime_ideals_over(L, typeof(p)[p])
    P1 = norm_1_generators(P)
    for x = P1
      push!(g, sum([e*preimage(mC, k) for (k, e) = x.fac]))
    end
  end

#  println("found $g")
  sC, msC = sub(C, collect(g))

  #think: might tbe better using quo, test for zero and a function
  #       that adds relations to a quo.
  if order(sC) != order(C)
    p = 2
    stable = 0
    while true
      while p in S
        p = next_prime(p)
      end
      change = false
      P1 =  norm_1_generators(prime_ideals_over(L, [p]))
      for x = P1
        y = sum([e*preimage(mC, k) for (k, e) = x.fac]) # move to log function(?)
        if !haspreimage(msC, y)[1]
          change = true
          push!(g, y)
          push!(S, p)
          sC, msC = sub(C, collect(g))
          if order(sC) == order(C)
#            println("groups are the same")
            break
          end
        end
      end
      p = next_prime(p)
      if change
        stable = 0
      else
        stable += 1
        if stable > 5
          break
        end
      end
    end
  else
#    println("groups are the same")
  end

  SP = prime_ideals_over(L, [x for x = S])

  U, mU = sunit_group_fac_elem(SP)
  u, mu = sunit_group_fac_elem(collect(S))

  h = elem_type(u)[]
  for i=1:ngens(U)
    x = mU(U[i])
    y = norm(x)
    push!(h, preimage(mu, y))
    #@assert norm(evaluate(x)) == evaluate(image(mu, h[end]))
  end
  s, ms = sub(u, h)
  mp = GrpAbFinGenMap(U, u, vcat([x.coeff for x=h]))

  fl, p = haspreimage(mp, preimage(mu, a))
  if fl
    return true, mU(p)
  else
    return false, FacElem(one(K))
  end
end


using Test, GenericLinearAlgebra, LinearAlgebra
using GenericLinearAlgebra.LAPACK2

@testset "LAPACK wrappers" begin

    n = 100

    T = SymTridiagonal(randn(n), randn(n - 1))
    vals, vecs = eigen(T)
    @testset "steqr" begin
        _vals, _vecs = LAPACK2.steqr!('N', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals

        _vals, _vecs = LAPACK2.steqr!('I', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals
        @test abs.(_vecs'vecs) ≈ Matrix(I, n, n)

        _vals, _vecs = LAPACK2.steqr!('V', copy(T.dv), copy(T.ev), Matrix{Float64}(I, n, n))
        @test vals ≈ _vals
        @test abs.(_vecs'vecs) ≈ Matrix(I, n, n)
    end

    @testset "sterf" begin
        _vals = LAPACK2.sterf!(copy(T.dv), copy(T.ev))
        @test _vals ≈ vals
    end

    @testset "stedc" begin
        _vals, _vecs = LAPACK2.stedc!('N', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals

        _vals, _vecs = LAPACK2.stedc!('I', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals
        @test abs.(_vecs'vecs) ≈ Matrix(I, n, n)

        _vals, _vecs = LAPACK2.stedc!('V', copy(T.dv), copy(T.ev), Matrix{Float64}(I, n,n))
        @test vals ≈ _vals
        @test abs.(_vecs'vecs) ≈ Matrix(I, n, n)
    end

    @testset "stemr" begin
        _vals, _vecs, __ = LAPACK2.stemr!('N', 'A', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals

        _vals, _vecs, __ = LAPACK2.stemr!('N', 'V', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals

        _vals, _vecs, __ = LAPACK2.stemr!('N', 'I', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals

        _vals, _vecs, __ = LAPACK2.stemr!('V', 'A', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals
        @test abs.(_vecs'vecs) ≈ Matrix(I, n, n)

        _vals, _vecs, __ = LAPACK2.stemr!('V', 'V', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals
        @test abs.(_vecs'vecs) ≈ Matrix(I, n, n)

        _vals, _vecs, __ = LAPACK2.stemr!('V', 'I', copy(T.dv), copy(T.ev))
        @test vals ≈ _vals
        @test abs.(_vecs'vecs) ≈ Matrix(I, n, n)
    end

    @testset "lahqr" begin
        T = Array(Tridiagonal(fill(0.1, 99), fill(0.0, 100), fill(40.0, 99)))
        _vals = sort(LAPACK2.lahqr!(copy(T))[1])
        @test _vals ≈ sort(real.(GenericLinearAlgebra._eigvals!(copy(T))))
        # LAPACK's multishift algorithm (the default) seems to be broken
        @test !(_vals ≈ sort(eigvals(T)))
    end
end


export JacobianConvexparameters,
    CharpolyCoeff,
    JacobianDeterminantConvexparameters

function JacobianConvexparameters(S::AbstractAlgebra.Ring,
                                  N::AbstractMatrix, Y::AbstractMatrix, E::AbstractMatrix,
                                  lpos::UnitRange{T}=1:size(E, 2),
                                  hpos::UnitRange{T}=.+(length(lpos), 1:size(N, 1))) where {T <: Integer}
    ## Migrating to AbstractAlgebra
    Nnemo = Nemo.matrix(S, N)
    Ynemo = Nemo.matrix(S, transpose(Y))
    ## Square matrix with h's diagonal
    H = Diagonal(S, Nemo.gens(S)[hpos])
    ## Diagonal with coordinates of E
    D = Diagonal(S, E * Nemo.gens(S)[lpos])
    ## Return the jacobian under convex parameters
    return Nnemo * D * Ynemo * H
end

function CharpolyCoeff(R::AbstractAlgebra.Ring, M::MatElem, ncoeff::Integer)
    P, x = Nemo.PolynomialRing(R, "x")
    # q = coeff(divexact(charpoly(P,Jhl),x^mindeg),0);
    return Nemo.coeff(charpoly(P, M), ncoeff)
end

function JacobianDeterminantConvexparameters(R::AbstractAlgebra.Ring,
                                             N::AbstractMatrix, Y::AbstractMatrix, E::AbstractMatrix,
                                             lpos::UnitRange{T}=1:size(E, 2),
                                             hpos::UnitRange{T}=.+(length(lpos), 1:size(N, 1)),
                                             mindeg::Integer=size(N, 1) - rank(Nemo.matrix(R, N))) where {T <: Integer}
    return CharpolyCoeff(R, JacobianConvexparameters(R, N, Y, E, lpos, hpos), mindeg)
end

function JacobianConvexparameters(H::AbstractAlgebra.Ring,
                                  L::AbstractAlgebra.Ring,
                                  N::AbstractMatrix, Y::AbstractMatrix, E::AbstractMatrix)
    ## Migrating to AbstractAlgebra
    Nnemo = Nemo.matrix(H, N)
    Ynemo = Nemo.matrix(H, transpose(Y))
    ## Square matrix with h's diagonal
    digH = Diagonal(H, Nemo.gens(H))
    ## Diagonal with coordinates of E
    digL = Diagonal(H, E * Nemo.gens(L))
    ## Return the jacobian under convex parameters
    return Nnemo * digL * Ynemo * digH
end

function JacobianDeterminantConvexparameters(H::AbstractAlgebra.Ring,
                                             L::AbstractAlgebra.Ring,
                                             N::AbstractMatrix, Y::AbstractMatrix, E::AbstractMatrix,
                                             mindeg::Integer=size(N, 1) - rank(Nemo.matrix(H, N)))
    return CharpolyCoeff(H, JacobianConvexparameters(H, L, N, Y, E), mindeg)
end


function test_adaptive_sketch(A, SketchType, r_true;
    tol=1e-6,
    ρ=1e-4,
    k_factor=0.9,
    condition_number=false,
)
    @testset "Adaptive Sketch Building" begin
        @testset "Norm" begin
            Anys_adapt = RP.adaptive_sketch(A, 2, SketchType; tol=10tol)
            @test round(r_true - rank(Anys_adapt) * 1/k_factor) == 0
            @test opnorm(A - Matrix(Anys_adapt)) <= tol
        end

        if condition_number
            @testset "Condition Number" begin
                Anys_adapt = RP.adaptive_sketch(A, 2, SketchType; tol=100tol, ρ=ρ, condition_number=true)
                @test round(r_true - rank(Anys_adapt) * 1/k_factor) == 0
                U, Λ = Anys_adapt.U, Anys_adapt.Λ
                M = Symmetric( (Λ[end] + ρ)*(U*inv(Λ + ρ*I)*U' + (I - U*U')) * (A + ρ*I) )
                @test (eigmax(M))/(eigmin(M) + ρ) - 1 <= tol
            end
        end
    end

    return nothing
end

@testset "Nystrom Sketch" begin
    # Data
    Random.seed!(0)
    n, r_true = 500, 100
    A = randn(n, r_true)
    A = A*A'

    r = round(Int, r_true * 1.2)
    k = round(Int, 0.9*r)
    λ = eigvals(A; sortby=x->-x)
    Anys = RP.NystromSketch(A, k, r)
    @test size(Anys, 1) == n && size(Anys, 2) == n
    @test rank(Anys) == k
    @test λ[1:k] ≈ svdvals(Anys)

    @test λ[1:k] ≈ eigvals(Anys)
    @test ≈(RP.estimate_norm_E(A, Anys; q=20), opnorm(A - Matrix(Anys)); rtol=1e-2)

    k, r = k ÷ 2, r ÷ 2
    Anys = RP.NystromSketch(A, k, r)
    @test ≈(RP.estimate_norm_E(A, Anys; q=50), opnorm(A - Matrix(Anys)); rtol=2e-1)

    x = randn(n)
    y = Anys * x
    @test y ≈ Matrix(Anys) * x
    z = zeros(n)

    test_adaptive_sketch(A, RP.NystromSketch, 128; k_factor=0.9, ρ=1e-4, tol=1e-6, condition_number=true)
end

@testset "EigenSketch" begin
    # Data
    Random.seed!(0)
    n, r1, r2 = 500, 60, 40
    r_true = r1+r2
    AAT(n, r) = randn(n, r) |> A->A*A'
    A = AAT(n, r1) - AAT(n, r2)

    r = round(Int, r_true * 1.2)
    k = round(Int, 0.9*r)
    A_sketch = RP.EigenSketch(A, k, r; check=true, q=5)
    @test size(A_sketch, 1) == n && size(A_sketch, 2) == n
    @test rank(A_sketch) == k
    λ = eigvals(A; sortby=x->-x) |> x->x[abs.(x) .≥ 1e-8]
    λhat =  eigvals(A_sketch) |> x->x[abs.(x) .≥ 1e-8]

    @test λ ≈ λhat
    @test RP.estimate_norm_E(A, A_sketch; q=20) <= 1e-8 && opnorm(A - Matrix(A_sketch)) <= 1e-8

    k, r = k ÷ 2, r ÷ 2
    A_sketch = RP.EigenSketch(A, k, r; q=100)
    @test ≈(RP.estimate_norm_E(A, A_sketch; q=100), opnorm(A - Matrix(A_sketch)); rtol=2e-1)

    x = randn(n)
    y = A_sketch * x
    @test y ≈ Matrix(A_sketch) * x
    z = zeros(n)

    test_adaptive_sketch(A, RP.EigenSketch, 128; k_factor=0.9, tol=1e-6)
end

@testset "Randomized SVD" begin
    # Data
    Random.seed!(0)
    n, r1, r2 = 500, 60, 40
    r = r1+r2
    AAT(n, r) = randn(n, r) |> A->A*A'
    A = AAT(n, r1) - AAT(n, r2)
    B = randn(n, n÷2) |> svd |> x->(x.U[:, 1:r] * Diagonal(x.S[1:r]) * x.Vt[1:r, :])
    B = B*B'*B

    for M in [A, B]
        r = round(Int, r1 + r2 * 1.2)
        k = round(Int, 0.9*r)
        Mhat = RP.RandomizedSVD(M, k, r; q=5)
        @test size(Mhat, 1) == size(M, 1) && size(Mhat, 2) == size(M, 2)
        @test rank(Mhat) == k
        @test svdvals(M)[1:k] ≈ svdvals(Mhat)
        @test ≈(RP.estimate_norm_E(M, Mhat; q=20), opnorm(M - Matrix(Mhat)); rtol=1e-1)

        k, r = k ÷ 2, r ÷ 2
        Mhat = RP.RandomizedSVD(M, k, r; q=5)
        @test ≈(RP.estimate_norm_E(M, Mhat; q=50), opnorm(M - Matrix(Mhat)); rtol=1e-1)

        x = randn(size(Mhat, 2))
        y = Mhat * x
        @test y ≈ Matrix(Mhat) * x
    end

    test_adaptive_sketch(A, RP.RandomizedSVD, 128; k_factor=0.9, tol=1e-6)

end


using SpinMC

a1 = (1.0, 0.0, 0.0)
a2 = (0.0, 1.0, 0.0)
a3 = (0.0, 0.0, 1.0)
uc = UnitCell(a1,a2,a3)

b = addBasisSite!(uc, (0.0, 0.0, 0.0))
M = [-1.0 0.0 0.0; 0.0 -1.0 0.0; 0.0 0.0 -1.0]
addInteraction!(uc, b, b, M, (1, 0, 0))
addInteraction!(uc, b, b, M, (0, 1, 0))
addInteraction!(uc, b, b, M, (0, 0, 1))

L = (8, 8, 8)
lattice = Lattice(uc, L)
thermalizationSweeps = 10000
measurementSweeps = 100000

tmin = 0.1
tmax = 10.0
N = 32
temperature = zeros(N)
heat = zeros(N)
dheat = zeros(N)
magnetization = zeros(N)
dmagnetization = zeros(N)
for i in 1:N
    println("Performing MC simulation ", i, "/", N)
    beta = 1.0 / (tmax * (tmin / tmax)^((i-1)/(N-1)))
    m = MonteCarlo(lattice, beta, thermalizationSweeps, measurementSweeps, reportInterval=20000)
    run!(m)
    temperature[i] = 1.0 / beta

    # Calculate specific heat
    # (Note that if file output is activated in the run!() call, the specific heat is automatically included.)
    c(e) = beta * beta * (e[2] - e[1] * e[1]) * length(m.lattice)
    ∇c(e) = [-2.0 * beta * beta * e[1] * length(m.lattice), beta * beta * length(m.lattice)]
    heat[i] = mean(m.observables.energy, c)
    dheat[i] = std_error(m.observables.energy, ∇c)
    # Calculate magnetization
    magnetization[i] = mean(m.observables.magnetization)
    dmagnetization[i] = std_error(m.observables.magnetization)
end

using Plots
display(plot(temperature, heat, yerror=dheat, xaxis=:log, yrange=(0.0,2.5), xrange=(0.1,10.0), xlabel="temperature", ylabel="specific heat", label="8*8*8 cubic lattice"))
display(plot(temperature, magnetization, yerror=dmagnetization, xaxis=:log, yrange=(0.0,1.0), xrange=(0.1,10.0), xlabel="temperature", ylabel="magnetization", label="8*8*8 cubic lattice"))


mutable struct Parameters{T}
    diam::T
    rmax::T
    ρ::T
    mr::Integer
    nρ::Integer
    δρ::T
    ktemp::T
end

function Parameters(diam, rmax, ρ, mr::Integer, nr, ktemp)
    (d, rm, r, n) = promote(diam, rmax, ρ, nr)

    return Parameters{typeof(d)}(d, rm, r, mr, n, zero(typeof(d)), ktemp)
end

struct Structure{T}
    r::T
    q::T
    pot::Potential
end

function Structure(p::Parameters, pot::Potential)
    dr = p.rmax / p.mr
    dq = π / p.rmax
    r = range(0.0, (p.mr - 1) * dr; length=p.mr)
    q = range(0.0, (p.mr - 1) * dq; length=p.mr)

    return Structure{typeof(r)}(r, q, pot)
end

struct Result{T<:AbstractArray}
    r::T
    q::T
    gr::T
    cr::T
    sq::T
end

function Result()
    vecs = [zeros(2) for _ in 1:5]
    return Result{Vector{Float64}}(vecs...)
end

function Result(cr, gam, st::Structure, params::Parameters)
    # First, we deal with the RDF
    gr = @. gam[2:end] + cr[2:end] + 1.0
    result_cr = copy(cr)

    # Next, we deal with the structure factor
    fits = continuous!(cr, st.r, params.diam)

    r, p = make_fft_plan(cr)
    ck = fft_oz(cr, r, params.rmax, params.mr, p)
    remove_continuous!(ck, st.q, params.diam, fits)

    sq_inv = @. 1.0 - params.ρ * ck[2:end]

    # Also save the new structure information
    new_r = collect(st.r[2:end])
    new_q = collect(st.q[2:end])

    return Result{typeof(gr)}(new_r, new_q, gr, result_cr, 1.0 ./ sq_inv)
end

struct Interaction
    params::Parameters
    structure::Structure
    bridge::Closure
end


module interaction

using ..system
using LinearAlgebra

export GreenTensor,G_ij,Gamma_ij,Omega_ij

"""
    interaction.F(ri::Vector, rj::Vector, µi::Vector, µj::Vector)

General F function for arbitrary positions and dipole orientations.

Arguments:
* ri: Position of first spin
* rj: Position of second spin
* µi: Dipole orientation of first spin.
* µj: Dipole orientation of second spin.
"""
function F(ri::Vector, rj::Vector, µi::Vector, µj::Vector)
    rij = ri - rj
    rij_norm = norm(rij)
    rijn = rij./rij_norm
    μi_ = normalize(μi)
    μj_ = normalize(μj)
    T = float(promote_type(eltype(rij),eltype(μi_),eltype(μj_)))
    if rij_norm == 0
        T(2/3.)
    else
        ξ = 2π*rij_norm
        T(dot(µi_, µj_)*(sin(ξ)/ξ + cos(ξ)/ξ^2 - sin(ξ)/ξ^3) + dot(µi_, rijn)*dot(rijn, µj_)*(-sin(ξ)/ξ - 3*cos(ξ)/ξ^2 + 3*sin(ξ)/ξ^3))
    end
end

"""
    interaction.G(ri::Vector, rj::Vector, µi::Vector, µj::Vector)

General G function for arbitrary positions and dipole orientations.

Arguments:
* ri: Position of first spin
* rj: Position of second spin
* µi: Dipole orientation of first spin.
* µj: Dipole orientation of second spin.
"""
function G(ri::Vector, rj::Vector, µi::Vector, µj::Vector)
    rij = ri - rj
    rij_norm = norm(rij)
    rijn = rij./rij_norm
    μi_ = normalize(μi)
    μj_ = normalize(μj)
    T = float(promote_type(eltype(rij),eltype(μi_),eltype(μj_)))
    if rij_norm == 0
        zero(T)
    else
        ξ = 2π*rij_norm
        T(dot(µi_, µj_)*(-cos(ξ)/ξ + sin(ξ)/ξ^2 + cos(ξ)/ξ^3) + dot(µi_, rijn)*dot(rijn, µj_)*(cos(ξ)/ξ - 3*sin(ξ)/ξ^2 - 3*cos(ξ)/ξ^3))
    end
end

"""
    interaction.Omega(ri::Vector, rj::Vector, µi::Vector, µj::Vector, γi::Real=1, γj::Real=1)

Arguments:
* ri: Position of first spin
* rj: Position of second spin
* µi: Dipole orientation of first spin.
* µj: Dipole orientation of second spin.
* γi: Decay rate of first spin.
* γj: Decay rate of second spin.

Note that the dipole moments `μi` and `μj` are normalized internally. To account
for dipole moments with different lengths you need to scale the decay rates
`γi` and `γj`, respectively.
"""
function Omega(ri::Vector, rj::Vector, µi::Vector, µj::Vector, γi::Real=1, γj::Real=1)
    return 0.75*sqrt(γi*γj)*G(ri, rj, µi, µj)
end

"""
    interaction.Gamma(ri::Vector, rj::Vector, µi::Vector, µj::Vector, γi::Real=1, γj::Real=1)

Arguments:
* ri: Position of first spin
* rj: Position of second spin
* µi: Dipole orientation of first spin.
* µj: Dipole orientation of second spin.
* γi: Decay rate of first spin.
* γj: Decay rate of second spin.

Note that the dipole moments `μi` and `μj` are normalized internally. To account
for dipole moments with different lengths you need to scale the decay rates
`γi` and `γj`, respectively.
"""
function Gamma(ri::Vector, rj::Vector, µi::Vector, µj::Vector, γi::Real=1, γj::Real=1)
    return 1.5*sqrt(γi*γj)*F(ri, rj, µi, µj)
end

"""
    interaction.OmegaMatrix(S::SpinCollection)

Matrix of the dipole-dipole interaction for a given SpinCollection.
"""
function OmegaMatrix(S::system.SpinCollection)
    spins = S.spins
    mu = S.polarizations
    gamma = S.gammas
    N = length(spins)
    Ω = zeros(Float64, N, N)
    for i=1:N, j=1:N
        if i==j
            continue
        end
        Ω[i,j] = interaction.Omega(spins[i].position, spins[j].position, mu[i], mu[j], gamma[i], gamma[j])
    end
    return Ω
end

"""
    interaction.GammaMatrix(S::SpinCollection)

Matrix of the collective decay rate for a given SpinCollection.
"""
function GammaMatrix(S::system.SpinCollection)
    spins = S.spins
    mu = S.polarizations
    gamma = S.gammas
    N = length(spins)
    return [
        interaction.Gamma(spins[i].position, spins[j].position, mu[i], mu[j], gamma[i], gamma[j])
        for i=1:N, j=1:N
    ]
end

"""
    GreenTensor(r::Vector, k::Number=2π)

Calculate the Green's Tensor at position r for wave number k defined by

```math
G = e^{ikr}\\Big[\\left(\\frac{1}{kr} + \\frac{i}{(kr)^2} - \\frac{1}{(kr)^3}\\right)*I -
    \\textbf{r}\\textbrf{r}^T\\left(\\frac{1}{kr} + \\frac{3i}{(kr)^2} - \\frac{3}{(kr)^3}\\right)\\Big]
```

Choosing `k=2π` corresponds to the position `r` being given in units of the
wavelength associated with the dipole transition.

Returns a 3×3 complex Matrix.
"""
function GreenTensor(r::Vector{<:Number},k::Real=2π)
    n = norm(r)
    rn = r./n
    return exp(im*k*n)*(
        (1/(k*n) + im/(k*n)^2 - 1/(k*n)^3).*Matrix(I,3,3) +
        -(1/(k*n) + 3im/(k*n)^2 - 3/(k*n)^3).*(rn*rn')
    )
end

"""
    Gamma_ij(r1::Vector,r2::Vector,μ1::Vector,μ2::Vector;gamma1=1,gamma2=1,k0=2π)

From the field overlap with the Green's tensor, compute the mutual decay rate as

```math
Γ_{ij} = \\frac{3}{2} μᵢ* ⋅ ℑ(G) ⋅ μⱼ.
```

"""
function Gamma_ij(r1::Vector{<:Real},r2::Vector{<:Real},μ1::Vector{<:Number},μ2::Vector{<:Number};k0=2π,kwargs...)
    T = float(promote_type(eltype(r1),eltype(r2),eltype(μ1),eltype(μ2),typeof(k0)))
    if (r1 == r2) && (μ1 == μ2)
        return one(T)
    elseif (r1 == r2) && (μ1 != μ2)
        return zero(T)
    else
        G_im = imag(GreenTensor(r1 - r2, k0))
        return 1.5*dot(μ1, G_im, μ2)::T
    end
end

"""
    Omega_ij(r1::Vector,r2::Vector,μ1::Vector,μ2::Vector;gamma1=1,gamma2=1,k0=2π)

From the field overlap with the Green's tensor, compute the energy shifts as

```math
Ω_{ij} = -\\frac{3}{4} μᵢ* ⋅ ℜ(G) ⋅ μⱼ.
```

"""
function Omega_ij(r1::Vector{<:Real},r2::Vector{<:Real},μ1::Vector{<:Number},μ2::Vector{<:Number};k0=2π,kwargs...)
    T = float(promote_type(eltype(r1),eltype(r2),eltype(μ1),eltype(μ2),typeof(k0)))
    if r1 == r2
        return zero(T)
    else
        G_re = real(GreenTensor(r1 - r2, k0))
        return -0.75*dot(μ1,G_re,μ2)::T
    end
end

end # module


#!/usr/bin/julia
function value(list)
    result = 0
    for v in list
        result = result * 2
        result = result + v
    end
    return result
end

input_file = ARGS[1]
open(input_file) do io
    lines = readlines(io)
    sums = sum(map.(x -> parse(Int64,x),map(collect, lines)))
    half = length(lines) / 2
    bits = map(v -> v > half ? 1 : 0, sums)
    gamma = value(bits)
    epsilon = value([1-v for v in bits])
    println(gamma * epsilon)
end


using StableRNGs

"""
Function to initialise the parameters of the desired network.
"""
function initialise_model(layer_dims::Vector{Int}, seed::Any)
    params = Dict()

    for l = 2:length(layer_dims)
        params[string("W_", (l-1))] = ( rand( StableRNG(seed), layer_dims[l], layer_dims[l-1] ) .- 0.5) * 200#* sqrt(layer_dims[l-1])
        params[string("b_", (l-1))] = zeros(layer_dims[l], 1)
    end # for

    return params
end


using DocumentDistances

sdd = SinkhornDocumentDistance()

# From Kusner's WMD paper they used this illustrative example of two
# quite similar sentences (documents) which are similar semantically
# but not syntactically.
d1 = ["obama", "speaks", "media", "illinois"]
d2 = ["president", "greets", "press", "chicago"]
d12 = evaluate(sdd, d1, d2)

# To compare to something let's take a different document which we expect
# to have a larger distance to the ones above than they have to each other.
d3 = ["lawyer", "tanks", "car", "africa"]

d13 = evaluate(sdd, d1, d3)
d23 = evaluate(sdd, d2, d3)

@assert d12 < d13
@assert d12 < d23


# routines for creating interpolatable forms of grid data

export interpolated_history

"""
    interpolated_history(v::History,tr)

Returns a function that interpolates the history `v`, using the time specified
by time range `tr`. Uses cubic spline interpolation. If `v` is a periodic history,
then the resulting function is also periodic.
"""
function interpolated_history(v::History{T,PeriodicHistory},tr::AbstractRange) where {T}
  length(tr) == length(v)+1 || error("Incommensurate lengths of history and range")
  return CubicSplineInterpolation(tr,Base.unsafe_view(v,1:length(tr)),extrapolation_bc=Periodic())
end

function interpolated_history(v::History{T,RegularHistory},tr::AbstractRange) where {T}
  length(tr) == length(v) || error("Incommensurate lengths of history and range")
  return CubicSplineInterpolation(tr,v,extrapolation_bc=Flat())
end

"""
    interpolated_history(v::History,xr,yr,tr)

Returns a function that interpolates the history `v` both spatially and temporally,
using the x and y ranges specified by `xr` and `yr` and time specified by time range
`tr`. Uses cubic spline interpolation in all directions. If `v` is a periodic history,
then the resulting function is also periodic in time.
"""
function interpolated_history(v::History{T,H},xr::AbstractRange,yr::AbstractRange,tr::AbstractRange) where {T,H<:HistoryType}
  # v is a PeriodicHistory, so need to add first element to last here
  nx, ny = size(v[1])

  extralen = (H == PeriodicHistory) ? 1 : 0

  (length(xr) == nx && length(yr) == ny && length(tr) == length(v)+extralen) || error("Incommensurate lengths of history and range")

  vhist = zeros(nx,ny,length(v)+extralen)
  for i in 1:size(vhist,3)
    vhist[:,:,i] = v[i]
  end

  bctype = (H == PeriodicHistory) ? Periodic() : Flat()

  return CubicSplineInterpolation((xr, yr ,tr), vhist, extrapolation_bc = (Flat(),Flat(),bctype))

end


# comparing EQ, MaternP, NN kernel with energetic norm for outer loop data
using LinearAlgebra
using SARA
using SARA: iSARA, random_sampling
using CovarianceFunctions
const Kernel = CovarianceFunctions
using GaussianDistributions
using Base.Threads

# generalized coefficient of determination, where errors are weighted by their standard deviation
# McFadden (in the context of logistic regression), (log?) likelihood ratio index
# x is data, y is model prediction, σ is expected noise for each data point
function likelihood_ratio(x::AbstractVector, y::AbstractVector, σ::AbstractVector = ones(length(x)))
    δ = (x .- mean(x)) ./ σ
    tss = sum(abs2, δ)
    δ = (x .- y) ./ σ
    ess = sum(abs2, δ)
    r2 = 1 - ess / tss
end

################# setup kernels for comparison ##############
# θ are stripe-dependent rescaling parameters
# LSA-bias is encoded in last parameter tuple
# l is length scale for all kernels
get_kernels(l) = get_kernels(l...)
function get_kernels(T_length::Real, τ_length::Real)
    # list of base kernels
    k1 = Kernel.EQ()
    k2 = Kernel.MaternP(2)
    k3 = Kernel.NN(1000.) # parameter governs size of bias coefficients in neural network, put proportional to temperature domain
    base_kernels = [k1, k2, k3]
    # equip base kernels with energetic norm
    # putlength(k) = Kernel.Lengthscale(k, l)
    # A = Diagonal(inv.([T_length, τ_length]).^2)
    # putenergy(k) = Kernel.Energetic(k, A)
    # kernels = putenergy.(base_kernels)
    scaling = [T_length, τ_length]
    scale(k) = (x, y) -> k(x./scaling, y./scaling)
    kernels = scale.(base_kernels)
    return kernels
end

############### generate sampling paths with different kernels
function sampling_path(G::Gaussian, x::AbstractMatrix,
                       y::AbstractMatrix,  σ²::AbstractMatrix, samplesizes)
    C = G
    nsample = length(samplesizes)
    rmse, r2 = zeros(nsample), zeros(nsample)
    indices = zeros(Int, 0)
    n, m = size(x)
    helper(x, i) = @views vec(x[:, i]) # selects the columns indexed by i, and vectorizes the result
    for (i, s) in enumerate(samplesizes) # random sampling
        indices = rand(1:m, s) # choose s random columns
        sort!(indices)
        xi, yi, σ²i = helper(x, indices), helper(y, indices), helper(σ², indices)
        C = conditional(G, xi, Gaussian(yi, σ²i), tol = 1e-6)
        μ = mean(C, vec(x))
        δ = vec(y) - μ
        rmse[i] = sqrt(mean(δ.^2 ./ vec(σ²))) # weighted root mean squared deviation
        r2[i] = likelihood_ratio(vec(y), μ, vec(σ²))
    end
    return rmse, r2
end

# using random sampling
function outer_kernel_benchmark!(rmse, r2, kernels, x, y, σ², samplesizes)
    nexp = size(rmse, 1)
    for (j, k) in enumerate(kernels)
        println("kernel $j of $(length(kernels))")
        G = Gaussian(k)
        @threads for i in 1:nexp
            d, r = sampling_path(G, x, y, σ², samplesizes)
            rmse[i, :, j] = d
            r2[i, :, j] = r
        end
    end
    return rmse
end

function energetic_kernel_benchmark(x, y, σ², nexp, nkernel, samplesizes, samplesperstripe)
    nsample = length(samplesizes) # number of different sample sizes [1, ..., nsample]
    # range of plausible length scales
    T_length = range(20, 40, step = 10)
    τ_length = range(.2, .4, step = .1)
    nlength = length(T_length) * length(τ_length)
    rmse = zeros(nexp, nsample, nkernel, length(T_length), length(τ_length))
    r2 = zeros(nexp, nsample, nkernel, length(T_length), length(τ_length))
    # first, calculate with energetic input norm
    for (i, T_l) in enumerate(T_length)
        for (j, τ_l) in enumerate(τ_length)
            println("lengthscale $(length(τ_length)*(i-1) + j) of $nlength")
            kernels = get_kernels((T_l, τ_l))
            rmse_ij = @view rmse[:, :, :, i, j]
            r2_ij = @view r2[:, :, :, i, j]
            outer_kernel_benchmark!(rmse_ij, r2_ij, kernels, x, y, σ², samplesizes)
        end
    end

    if true
        f = h5open("outer_kernel_benchmark_energetic.h5", "w")
        f["weighted rmse"] = rmse
        f["likelihood ratio"] = r2
        f["temperature lengthscales"] = collect(T_length) # in resultss
        f["dwelltime lengthscales"] = collect(τ_length)
        f["nexp"] = nexp
        f["sample sizes"] = collect(samplesizes)
        f["nlength"] = nlength
        f["nkernel"] = nkernel
        f["kernel names"] = ["EQ", "Matern_2", "NN"]
        f["dimensions"] = ["nexp", "nsample", "nkernel", "temperature lengthscales", "dwelltime lengthscales"]
        f["samples per stripe"] = samplesperstripe
        close(f)
    end
    return rmse, r2
end

function uniform_kernel_benchmark(x, y, σ², nexp, nkernel, samplesizes, samplesperstripe)
    nsample = length(samplesizes) # number of different sample sizes [1, ..., nsample]
    # range of plausible length scales
    lengthscales = [.25, .5, 1., 2., 4., 8., 16., 32.]
    rmse = zeros(nexp, nsample, nkernel, length(lengthscales))
    r2 = zeros(nexp, nsample, nkernel, length(lengthscales))
    for (i, l) in enumerate(lengthscales)
        println("lengthscale $i of $(length(lengthscales))")
        kernels = get_kernels((l, l))
        rmse_i = @view rmse[:, :, :, i]
        r2_i = @view r2[:, :, :, i]
        outer_kernel_benchmark!(rmse_i, r2_i, kernels, x, y, σ², samplesizes)
    end

    if true
        f = h5open("outer_kernel_benchmark_uniform.h5", "w")
        f["weighted rmse"] = rmse
        f["likelihood ratio"] = r2
        f["lengthscales"] = collect(lengthscales) # in resultss
        f["nexp"] = nexp
        f["sample sizes"] = collect(samplesizes)
        f["nlength"] = length(lengthscales)
        f["nkernel"] = nkernel
        f["kernel names"] = ["EQ", "Matern_2", "NN"]
        f["dimensions"] = ["nexp", "nsample", "nkernel", "lengthscales"]
        f["samples per stripe"] = samplesperstripe
        close(f)
    end
    return rmse, r2
end

function outer_kernel_benchmark_driver(x::AbstractMatrix, y::AbstractMatrix, σ²::AbstractMatrix)
    # subsample observations per
    subsample(x, k::Int) = @view x[end+1-k:end, :]
    samplesperstripe = 8
    x = subsample(x, samplesperstripe)
    y = subsample(y, samplesperstripe)
    σ² = subsample(σ², samplesperstripe)

    # standardize global data
    y = y .- mean(y)
    y_normalization = std(y)
    y /= y_normalization
    σ² /= y_normalization^2 # noise variance

    # experimental parameters
    nexp = 1 # number of independent experiments per nsample
    nkernel = 3 # number of kernels

    samplemax = 64
    samplestep = 32 # number of sample
    samplesizes = range(32, samplemax, step = samplestep)
    energetic_kernel_benchmark(x, y, σ², nexp, nkernel, samplesizes, samplesperstripe)
    uniform_kernel_benchmark(x, y, σ², nexp, nkernel, samplesizes, samplesperstripe)
end

# load outer loop data
using HDF5
path = "/Users/sebastianament/Documents/SEA/XRD Analysis/SARA/Bi2O3_19F44_01/"
# path = "NatCom2020/"
# file = "Bi2O3_19F44_01_outer_loop_data.h5" # without input noise
file = "Bi2O3_19F44_01_outer_loop_input_noise_data.h5"
f = h5open(path * file, "r")
dwelltimes = read(f, "dwelltimes")
temperatures = read(f, "temperatures")
gradients = read(f, "gradients")
var_gradients = read(f, "var(gradients)")
close(f)

# pre-process data
x = tuple.(temperatures, dwelltimes)
y = gradients
σ² = var_gradients # uncertainty in y

rmse, r2 = outer_kernel_benchmark_driver(x, y, σ²)

# mrmse = reshape(mean(rmse, dims = 1), size(rmse)[2:4]..., :)
# mr2 = reshape(mean(r2, dims = 1), size(r2)[2:4]..., :)


#===============================================================================
            Some simple online statistics for the Markov chain
===============================================================================#

"""
    mutable struct GenericChainStats{T} <: ChainStats
        rolling_ar::Vector{Vector{Float64}}
        roll_window::Int64
        mean::Vector{T}
        cov::Matrix{T}
        N::Int64
    end

Simple online statistics for the Markov chain.
"""
mutable struct GenericChainStats{T} <: ChainStats
    rolling_ar::Vector{Vector{Float64}}
    roll_window::Int64
    mean::Vector{T}
    cov::Matrix{T}
    N::Int64

    function GenericChainStats(
            state::Vector{T},
            updt_len,
            num_mcmc_steps,
            roll_window=100
        ) where T
        new{T}(
            [zeros(Float64, updt_len) for _=1:num_mcmc_steps],
            roll_window,
            zero(state),
            zeros(T, (length(state), length(state))),
            1
        )
    end

end

function update_stats!(
        global_ws::GlobalWorkspace,
        local_ws::LocalWorkspace,
        step
    )
    cs = global_ws.sub_ws.stats
    θ = global_ws.sub_ws.state
    old_sum_sq = (cs.N-1)/cs.N * cs.cov + cs.mean * cs.mean'
    cs.mean .= cs.mean .* (cs.N/(cs.N+1)) .+ θ ./ (cs.N+1)
    new_sum_sq = old_sum_sq + (θ * θ')/cs.N
    cs.cov .= new_sum_sq - (cs.N+1)/cs.N*(cs.mean * cs.mean')

    accepted_now = local_ws.acceptance_history[step.mcmciter]
    ra_prev = cs.rolling_ar[max(1,step.mcmciter-1)][step.pidx]
    accepted_outside_window = (
        step.mcmciter > cs.roll_window ?
        local_ws.acceptance_history[step.mcmciter-cs.roll_window] :
        false
    )

    cs.rolling_ar[step.mcmciter][step.pidx] = (
        ra_prev * cs.roll_window
        + (accepted_now - accepted_outside_window)
    )  / min(cs.roll_window, cs.N)
    cs.N += 1
end


@testset "Accuracy" begin

m = 5
σ = 2.0

@testset "High-level NFFT" begin
  @info "High-level NFFT"
  N = (32,32)
  eps =  1e-3
  D = length(N)

  M = prod(N)
  x = rand(Float64,D,M) .- 0.5

  fHat = rand(Float64,M) + rand(Float64,M)*im
  f = ndft_adjoint(x, N, fHat,)
  fApprox = nfft_adjoint(x, N, fHat, reltol=1e-9)
  e = norm(f[:] - fApprox[:]) / norm(f[:])
  @debug "error adjoint nfft "  e
  @test e < eps

  gHat = ndft(x, f)
  gHatApprox = nfft(x, f, reltol=1e-9)
  e = norm(gHat[:] - gHatApprox[:]) / norm(gHat[:])
  @debug "error nfft "  e
  @test e < eps

  f = ndft_adjoint(x, N, ComplexF32.(fHat))
  fApprox = nfft_adjoint(x, N, ComplexF32.(fHat))
  e = norm(f[:] - fApprox[:]) / norm(f[:])
  @debug "error adjoint nfft "  e
  @test e < eps

  gHat = ndft(x, ComplexF32.(f))
  gHatApprox = nfft(x, ComplexF32.(f))
  e = norm(gHat[:] - gHatApprox[:]) / norm(gHat[:])
  @debug "error nfft "  e
  @test e < eps
end

@testset "NFFT in multiple dimensions" begin
    for (u,N) in enumerate([(256,), (30,32), (10,12,14), (6,6,6,6)])
      for (pre, storeDeconv, blocking) in zip([ NFFT.LINEAR, NFFT.LINEAR, NFFT.LINEAR, NFFT.FULL, NFFT.TENSOR, NFFT.POLYNOMIAL, NFFT.POLYNOMIAL],
                                           [false, true, false, false, false, false, false],
                                           [true, true, false, false, true, true, false])
        eps = [1e-7, 1e-7, 1e-3, 1e-6, 1e-4]
        for (l,window) in enumerate([:kaiser_bessel, :cosh_type, :gauss, :kaiser_bessel_rev, :spline])
            D = length(N)
            @info "Testing $D, window=$window, pre=$pre, storeDeconv=$storeDeconv, block=$blocking"

            M = prod(N)
            x = rand(Float64,D,M) .- 0.5
            p = plan_nfft(x, N; m, σ, window, precompute = pre, storeDeconvolutionIdx = storeDeconv,
                         fftflags = FFTW.ESTIMATE, blocking)

            pNDFT = NDFTPlan(x, N)

            fHat = rand(Float64,M) + rand(Float64,M)*im
            f = adjoint(pNDFT) * fHat
            fApprox = adjoint(p) * fHat
            e = norm(f[:] - fApprox[:]) / norm(f[:])
            @debug "error adjoint nfft "  e
            @test e < eps[l]

            gHat = pNDFT * f
            gHatApprox = p * f
            e = norm(gHat[:] - gHatApprox[:]) / norm(gHat[:])
            @debug "error nfft "  e
            @test e < eps[l]
        end
      end
    end
end

@testset "Abstract sampling points" begin
  @info "Abstract sampling points"
    M, N = rand(100:2:200, 2)
    x = range(-0.4, stop=0.4, length=M)
    p = plan_nfft(x, N, fftflags = FFTW.ESTIMATE)
end

@testset "Directional NFFT $D dim" for D in 2:3 begin
    # NFFT along a specified dimension should give the same result as
    # running a 1D NFFT on every slice along that dimension
        eps = 1e-4
        N = tuple( 2*rand(4:8,D)... )
        M = prod(N)
        for d in 1:D
            @info "Testing in $D dimensions directional NFFT along dim=$d"
            x = rand(M) .- 0.5

            f = rand(ComplexF64,N)
            p_dir = plan_nfft(x, N, dims=d)
            fHat_dir = p_dir * f
            g_dir = adjoint(p_dir) * fHat_dir

            p = plan_nfft(x, N[d])
            fHat = similar(fHat_dir)
            g = similar(g_dir)

            sz = size(fHat)
            Rpre = CartesianIndices( sz[1:d-1] )
            Rpost = CartesianIndices( sz[d+1:end] )
            for Ipost in Rpost, Ipre in Rpre
                idx = [Ipre, :, Ipost]
                fview = f[idx...]
                fHat[idx...] = p * vec(fview)

                fHat_view = fHat_dir[idx...]
                g[idx...] = adjoint(p) * vec(fHat_view)
            end

            e = norm( fHat_dir[:] - fHat[:] )
            @test e < eps

            e = norm( g_dir[:] - g[:] ) / norm(g[:])
            @test e < eps
        end
    end
end

@testset "Directional NFFT $D dim" for D in 3:3 begin
  # NFFT along  specified dimensions (1:2 or 2:3) should give the same result as
  # running a 2D NFFT on every slice along that dimensions
      eps = 1e-4
      N = tuple( 2*rand(4:8,D)... )
      M = prod(N)
      for d in 1:(D-1)
          dims = d:(d+1)
          @info "Testing in $D dimensions directional NFFT along dim=$(dims)"
          x = rand(2,M) .- 0.5

          f = rand(ComplexF64, N)
          p_dir = plan_nfft(x, N, dims=dims)
          fHat_dir = p_dir * f
          g_dir = adjoint(p_dir) * fHat_dir

          p = plan_nfft(x, N[dims])
          fHat = similar(fHat_dir)
          g = similar(g_dir)

          Rpre = CartesianIndices( N[1:dims[1]-1] )
          Rpost = CartesianIndices( N[(dims[end]+1):end] )
          for Ipost in Rpost, Ipre in Rpre
              idxf = [Ipre, :, :, Ipost]
              idxfhat = [Ipre, :, Ipost]

              fview = f[idxf...]
              fHat[idxfhat...] = p * fview

              fHat_view = fHat_dir[idxfhat...]
              g[idxf...] = adjoint(p) * fHat_view
          end

          e = norm( fHat_dir[:] - fHat[:] )
          @test e < eps

          e = norm( g_dir[:] - g[:] ) / norm(g[:])
          @test e < eps
      end
  end
end

end


# Rounding
function round(x::Decimal; digits::Int=0, normal::Bool=false)
    shift = BigInt(digits) + x.q
    if shift > BigInt(0) || shift < x.q
        (normal) ? x : normalize(x, rounded=true)
    else
        c = Base.round(x.c / BigInt(10)^(-shift))
        d = Decimal(x.s, BigInt(c), x.q - shift)
        (normal) ? d : normalize(d, rounded=true)
    end
end

function trunc(x::Decimal; digits::Int=0, normal::Bool=false)
    shift = BigInt(digits) + x.q
    if shift > BigInt(0) || shift < x.q
        (normal) ? x : normalize(x, rounded=true)
    else
        c = Base.trunc(x.c / BigInt(10)^(-shift))
        d = Decimal(x.s, BigInt(c), x.q - shift)
        (normal) ? d : normalize(d, rounded=true)
    end
end


module PHcalc
using Optim

export Neutral,Acid, System, α,pHsolve, minimise, pHfast

mutable struct Acid
	ka::Vector{Float64}
	conc::Float64
	charge::Vector{Int64}
	pKA::Vector{Float64}
	_ka::Vector{Float64}
	function Acid(ka, conc, charge::Int64)
			if typeof(ka) <: Number
				ka=[ka]
			end
			_ka=copy(ka)
			sort!(ka, rev=:true)
			prepend!(_ka,1.)
			charge=charge:-1:charge-(length(ka)) |> collect
			# charge=LinRange(charge,-length(ka), length(ka))
			new(ka, conc,charge , -log10.(ka),_ka )
		end
end

mutable struct Neutral
	charge::Union{Vector{Int64},Int64}
	conc::Float64
end

struct System
	species
	function System(species...)
		new(species)
	end
end

function α(self::Acid,pH) # solo para un valor de pH, se podria pensar para un array de pHs.
		h3o=10.0^(-pH)
		power=length(self._ka)-1:-1:0
		h3o_power=(x->h3o^x).(power)
		Ka_prod=cumprod(self._ka)
		h3o_Ka=h3o_power.*(Ka_prod)
		h3o_Ka./sum(h3o_Ka)
	end

function α(self::Neutral,pH) # Devielve 1 porque no tiene disociacion, devuelvo en array para que se parezca a la otra implementación.
	[1]
end

function pHsolve(sys)
	function minimise(pH)
		h3o=10.0^(-pH)
     	oh = (10.0^(-14))/h3o
	 	x = (h3o - oh)
		for specie in sys.species
			x+=(specie.conc.*specie.charge.*α(specie,pH))|> sum
		end
	 	abs(x)
	end

	function phguess()
		phv=1:.1:14
		minimise.(phv) |> argmin |> x -> getindex(phv, x)
	end
	guess=phguess()
	optimize(x->minimise(first(x)), [guess], BFGS()).minimizer |> first
end

function minimise(sys::System,pH)
	h3o=10.0^(-pH)
     	oh = (10.0^(-14))/h3o
	 x = (h3o - oh)
	for specie in sys.species
		x+=(specie.conc.*specie.charge.*α(specie,pH))|> sum
	end
	 abs(x)
end

function pHfast(sys, precision=.01)
	function minimise(pH)
		h3o=10.0^(-pH)
     	oh = (10.0^(-14))/h3o
	 	x = (h3o - oh)
		for specie in sys.species
			x+=(specie.conc.*specie.charge.*α(specie,pH))|> sum
		end
	 	abs(x)
	end
	phv=1:precision:14
	minimise.(phv) |> argmin |> x -> getindex(phv,x)

end

end


struct FDSwell <: FDSpectralModel
    θ::Vector{Float64}
end

parameter(F::FDSwell) = F.θ

function sdf(F::FDSwell, ω)
    θ = F.θ
    α = θ[1]
    ωₚ = θ[2]
    κ = θ[3]
    sigω = sign(ω)
    ω = abs(ω)
    logbit = log(ω) - log(ωₚ)
    h = logbit/κ - κ
    sdf = α * exp(-0.5 * (h)^2)
    sdf *= (!isnan(sdf))
end

spreading(F::FDSwell, ω, ϕ) = WrappedGaussian(ϕ, F.θ[4], F.θ[5])

####

struct FDSwellJS <: FDSpectralModel
    θ::Vector{Float64}
end

parameter(F::FDSwellJS) = F.θ

spreading(F::FDSwellJS, ω, ϕ) = WrappedGaussian(ϕ, F.θ[5], F.θ[6])


# Computes classical Stein discrepancy based on the g and gradg coefficients
# for univariate target distributions:
#
#     max_{g} sum_i gobjcoeff(i) * g(x(i)) + gradgobjcoeff(i) * g'(x(i)))
#     subject to, for all z, y
#     |g(z)| <= c_{1} I[alpha < z < beta], |g'(z)| <= c_{2},
#     |g(z) - g(y)| <= c_{2} |z-y|,
#     |g'(z) - g'(y)| <= c_3 |z-y|,
#     |g(z) - g(y) - g'(y)(z - y)| <= (c_3/2) |z-y|^2,
#     |g(z) - g(y) - g'(z)(z - y)| <= (c_3/2) |z-y|^2
#    where (alpha,beta) is the support of target distribution.
#
# Despite this being an infinite dimensional feasible set, there
# is a set of finite dimensional constraints that can exactly
# recover the optimal solution g, g' on the points x_i.
#
# Args:
#  sample - SteinDiscrete object representing a sample
#  gobjcoefficients - a vector of coefficients for the objective for g
#  gradgobjcoefficients - a vector of coefficients for the objective for gprime
#  solver - optimization program solver supported by JuMP that can solve
#   a quadratically constrained quadratic program (QCQP)
#  operator - the operator being used for the objective
function affine_classical_discrepancy(sample::SteinDiscrete,
                                      gobjcoefficients::Array{Float64,1},
                                      gradgobjcoefficients::Array{Float64,1},
                                      solver::AbstractMathProgSolver,
                                      operator::AbstractString;
                                      c1::Float64=1.0,
                                      c2::Float64=1.0,
                                      c3::Float64=1.0,
                                      supportlowerbounds=[-Inf],
                                      supportupperbounds=[Inf])

    # get primary objects
    points = sample.support
    weights = vec(sample.weights)
    n = length(weights)
    ## reset the support bounds if needed
    if isa(supportlowerbounds, Number)
        supportlowerbounds = [supportlowerbounds]
    end
    if isa(supportupperbounds, Number)
        supportupperbounds = [supportupperbounds]
    end
    ## Find spanner edges (consecutive pairs)
    println("[Computing spanner edges]")
    edgetime = @elapsed edges = getspanneredges(points)
    println("\tusing $(size(edges,1)) of $(binomial(n,2)) edges")

    elapsed = @elapsed begin
        # start setting up the model
        m = Model(solver=solver)
        # Define variables and specify single variable bounds
        @variable(m, -c1 <= g[1:n] <= c1)
        @variable(m, -c2 <= gprime[1:n] <= c2)
        # Introduce classical Stein program slack variables
        @variable(m, -Inf <= tb[1:(n-1)] <= Inf)
        @variable(m, -Inf <= tu[1:(n-1)] <= Inf)
        # set objective
        gobj = g[1:n]'gobjcoefficients
        gprimeobj = gprime[1:n]'gradgobjcoefficients
        @objective(m, Max, gobj + gprimeobj)
        # add gprime constraints
        xdistances = points[2:n, 1] - points[1:(n-1), 1]
        for i=1:(n-1)
            @constraint(m, gprime[i] - gprime[i+1] <= c3*xdistances[i])
            @constraint(m, gprime[i] - gprime[i+1] >= -c3*xdistances[i])
        end
        # Introduce constraints tb >= Lb and tu >= Lu:
        # tb_i >= c3/2 * (x_{i+1} - x_i) - (g'(x_i) + g'(x_{i+1}))/2 - c2
        # tu_i >= c3/2 * (x_{i+1} - x_i) + (g'(x_i) + g'(x_{i+1}))/2 - c2
        slackoffsets = (c3/2) * xdistances - c2
        for i=1:(n-1)
            @constraint(m, tb[i] + 0.5 * gprime[i] + 0.5 * gprime[i+1] >= slackoffsets[i])
            @constraint(m, tu[i] - 0.5 * gprime[i] - 0.5 * gprime[i+1] >= slackoffsets[i])
        end
        # Introduce quadratic buffer constraints:
        # |g(x_i)| <= c_1 - (1/(2c_3)) g'(x_i)^2
        for i=1:n
            qexp = QuadExpr([gprime[i]], [gprime[i]], [1/(2*c3)], [g[i]]'[1.0])
            @constraint(m, qexp <= c1)
            qexp = QuadExpr([gprime[i]], [gprime[i]], [1/(2*c3)], [g[i]]'[-1.0])
            @constraint(m, qexp <= c1)
        end
        # Add sharp constraints linking g and g':
        # g(x_{i+1}) - g(x_i) + (g'(x_{i+1}) - g'(x_i))^2/(4c3) - (x_{i+1} - x_i)*(g'(x_i) + g'(x_{i+1}))/2 + (1/c3)(L_u)^2_+ <= (c3/4)(x_{i+1} - x_i)^2
        # g(x_i) - g(x_{i+1}) + (g'(x_{i+1}) - g'(x_i))^2/(4c3) + (x_{i+1} - x_i)*(g'(x_i) + g'(x_{i+1}))/2 + (1/c3)(L_b)^2_+ <= (c3/4)(x_{i+1} - x_i)^2
        scaledsquaredxdistances = (xdistances.^2) .* (c3/4)
        for i=1:(n-1)
            rexp = [g[i:(i+1)]; gprime[i:(i+1)]]'[1.0, -1.0,  xdistances[i]/2, xdistances[i]/2]
            qbexp = QuadExpr(
                [gprime[[i, i, i+1]]; tb[i]],
                [gprime[[i, i+1, i+1]]; tb[i]],
                [1.0, -2.0, 1.0, 4.0] ./ (4*c3),
                rexp
            )
            quexp = QuadExpr(
                [gprime[[i, i, i+1]]; tu[i]],
                [gprime[[i, i+1, i+1]]; tu[i]],
                [1.0, -2.0, 1.0, 4.0] ./ (4*c3),
                -1 * rexp
            )

            @constraint(m, qbexp <= scaledsquaredxdistances[i])
            @constraint(m, quexp <= scaledsquaredxdistances[i])
         end

        # get boundary of support
        alpha = supportlowerbounds[1]
        beta = supportupperbounds[1]
        # add support slack constraints if necessary
        if isfinite(alpha)
            xdistance = points[1,1] - alpha
            # |g(x_1)| <= c2 (x_1 - alpha)
            @constraint(m, g[1] <= c2*xdistance)
            @constraint(m, g[1] >= -c2*xdistance)
            # g(x_1) <= g'(x_1)*(x_1 - alpha) + (c3/2)*(x_1-alpha)^2 - 1/(2c3) max{g'(x_1) + c3(x_1 - alpha) - c2, 0}^2,
            # g(x_1) >= g'(x_1)*(x_1 - alpha) - (c3/2)*(x_1-alpha)^2 + 1/(2c3) max{-g'(x_1) + c3(x_1 - alpha) - c2, 0}^2
            @variable(m, -Inf <= alphab <= Inf)
            @variable(m, -Inf <= alphau <= Inf)
            @constraint(m, alphab - gprime[1] >= c3 * xdistance - c2)
            @constraint(m, alphau + gprime[1] >= c3 * xdistance - c2)
            rexp = [g[1], gprime[1]]'[1.0, -xdistance]
            qbexp = QuadExpr([alphab], [alphab], [1/(2*c3)], rexp)
            quexp = QuadExpr([alphau], [alphau], [1/(2*c3)], -1 * rexp)
            @constraint(m, qbexp <= (c3/2) * xdistance^2)
            @constraint(m, quexp <= (c3/2) * xdistance^2)
        end
        if isfinite(beta)
            xdistance = beta - points[end,1]
            # |g(x_n)| <= c2 (beta - x_n)
            @constraint(m, g[n] <= c2*xdistance)
            @constraint(m, g[n] >= -c2*xdistance)
            # -g(x_n) <= g'(x_n)*(beta-x_n) + (c3/2)*(beta-x_n)^2 - 1/(2c3) max{g'(x_n) + c3(beta-x_n) - c2, 0}^2
            # -g(x_n) >= g'(x_n)*(beta-x_n) - (c3/2)*(beta-x_n)^2 + 1/(2c3) max{-g'(x_n) + c3(beta-x_n) - c2, 0}^2
            @variable(m, -Inf <= betab <= Inf)
            @variable(m, -Inf <= betau <= Inf)
            @constraint(m, betab + gprime[n] >= c3 * xdistance - c2)
            @constraint(m, betau - gprime[n] >= c3 * xdistance - c2)
            rexp = [g[n], gprime[n]]'[1.0, xdistance]
            qbexp = QuadExpr([betab], [betab], [1/(2*c3)], rexp)
            quexp = QuadExpr([betau], [betau], [1/(2*c3)], -1 * rexp)
            @constraint(m, qbexp <= (c3/2) * xdistance^2)
            @constraint(m, quexp <= (c3/2) * xdistance^2)
    end
    # Solve the optimization program
    @time optimize!(m)
    print(termination_status(m))
    end
    solvetime = elapsed
    # Package the results
    objval = [getobjectivevalue(m)]
    gopt = getvalue(g)[1:n]
    gprimeopt = getvalue(gprime)[1:n]
    operatorgopt = (gopt .* gobjcoefficients) + (gprimeopt .* gradgobjcoefficients)

    AffineDiscrepancyResult(
        points,
        weights,
        edges,
        objval,
        gopt,
        gprimeopt,
        operatorgopt,
        edgetime,
        solvetime,
        operator
    )
end


# struct Simulation{DT <: Number, ET <: Equation{DT}}
#     Δt::DT
#     equ::ET
#     x::OffsetArray{DT,3,Array{DT,3}}

#     function Simulation(equ::ET, Δt::DT, nt::Int) where {DT, ET <: Equation{DT}}
#         x = OffsetArray(zeros(DT, size(equ.x₀)..., nt+1), axes(equ.x₀)..., 0:nt)
#         x[:,:,0] .= equ.x₀
#         new{DT,ET}(Δt, equ, x)
#     end
# end

struct Simulation{DT <: Number, FT <: Function}
    Δt::DT
    equ::Equation{DT,FT}
    x::OffsetArray{DT,3,Array{DT,3}}

    function Simulation(equ::Equation{DT,FT}, Δt::DT, nt::Int) where {DT, FT}
        x = OffsetArray(zeros(DT, size(equ.x₀)..., nt+1), axes(equ.x₀)..., 0:nt)
        x[:,:,0] .= equ.x₀
        new{DT,FT}(Δt, equ, x)
    end
end

ndims(sim::Simulation) = ndims(sim.equ)
nsamples(sim::Simulation) = nsamples(sim.equ)
ntimesteps(sim::Simulation) = lastindex(sim.x,3)

eachsample(sim::Simulation) = axes(sim.x,2)
eachtimestep(sim::Simulation) = axes(sim.x,3)[1:end]

function run!(sim::Simulation{DT}) where {DT}
    ẋ = zeros(DT, ndims(sim))
    for n in eachtimestep(sim)
        for i in eachsample(sim)
            sim.equ.f(ẋ, sim.x[:,i,n-1])
            sim.x[:,i,n] .= sim.x[:,i,n-1] .+ sim.Δt .* ẋ
        end
    end
    return sim.x
end


## -------------  Model

@with_kw mutable struct Model
    # Number of parameters
    K :: Int64 = 12
    # Number of households
    N :: Int64 = 10
    # Number of options
    M :: Int64 = 20
    # How often to simulate each type
    nSim :: Int64 = 200

    # Governs how deviation function shows intermediate results
    showFreq :: Int64 = 100
    currentIter :: Int64 = 0
    # No of classes when computing stats
    nEndowCl :: Int = 6
    nChoiceCl :: Int = 8
end

# No of calibrated params
function n_params(mdl :: Model)
    return mdl.K
end

## ---------  Model parameters

mutable struct ModelParams
    alphaV :: Vector{Float64}
    # betaV :: Vector{Float64}
    # gammaV :: Vector{Float64}
    # deltaV :: Vector{Float64}
end

function show_params(mp :: ModelParams)
    println("Alpha:  ");
    println(round.(mp.alphaV, digits = 3));
    # print("Beta:   "); println(round.(mp.betaV, digits = 3));
    # print("Gamma:  "); println(round.(mp.gammaV, digits = 3));
    # print("Delta:  "); println(round.(mp.deltaV, digits = 3));
end

## --------  Model stats

mutable struct ModelStats
    # Fraction of choices in each [choice class, endowment class, by endowment] bin
    frac_ceM  ::  Array{Float64, 2}
end

function n_tg_moments(ms :: ModelStats)
    return length(ms.frac_ceM)
end

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