# ZEROFPR
#
# PART OF BAZINGA.jl

using Base.Iterators
using ProximalAlgorithms.IterationTools
using LinearAlgebra: norm
using OptiMo
using Printf

###########################################################################
# Options
###########################################################################
struct ZEROFPR_options{R<:Real}
    alpha::R                        # ∈ (0, 1), e.g.: 0.95
    beta::R                         # ∈ (0, 1), e.g.: 0.5
    sigma::R                        # = beta * (1 - alpha) / 2 ∈ (0, 0.5)
    adaptive::Bool                  # enforce adaptive stepsize even if L is provided
    tau_min::R                      # minimum linesearch stepsize
    gamma_min::R                    # minimum forward-backward stepsize
    bigeps::R                       # small value used for tolerances
    grad_scaling::Bool              # perform gradient-based scaling?
    scaling_grad_max::R             # maximum gradient after scaling
    scaling_grad_target::Maybe{R}   # target value of gradient after scaling
    scaling_min_value::R            # minimum value of gradient-based scaling factor
    warnings::Bool                  # can print warnings?
end

###########################################################################
# Iterable
###########################################################################
struct ZEROFPR_iterable{R<:Real,Tx<:AbstractArray{R},TP,TH,TO<:ZEROFPR_options{R}}
    prob::TP                        # problem
    x0::Tx                          # initial point
    gamma::Maybe{R}                 # forward-backward stepsize
    H::TH                           # quasi-Newton object
    opts::TO                        # options
end

Base.IteratorSize(::Type{<:ZEROFPR_iterable}) = Base.IsInfinite()

###########################################################################
# State
###########################################################################
mutable struct ZEROFPR_state{R<:Real,Tx,TH}
    x::Tx                   # iterate
    f_x::R                  # (unscaled) value of smooth term
    grad_f_x::Tx            # (unscaled) gradient of smooth term
    gamma::R                # stepsize parameter of forward and backward steps
    y::Tx                   # forward point
    xbar::Tx                # forward-backward point
    g_xbar::R               # (unscaled) value of nonsmooth term (at xbar)
    res::Tx                 # fixed-point residual at iterate (= x - xbar)
    optim::R                # optimality measure
    H::TH                   # variable metric
    tau::Maybe{R}           # linesearch stepsize
    objec::R                # (unscaled) objective function at `xbar`
    scalobj::R              # scaling factor (> 0)
    fbe_x::R                # (scaled) forward-backward envelope at `x`
    # additional allocation
    grad_f_xbar::Tx
    xbarbar::Tx
    res_xbar::Tx
    d::Tx                   # search direction at `xbar`
    xbar_old::Maybe{Tx}
    res_xbar_old::Maybe{Tx}
end

###########################################################################
# Quadratic model
###########################################################################
f_model(state::ZEROFPR_state) = f_model(
    state.scalobj * state.f_x,
    state.scalobj .* state.grad_f_x,
    state.res,
    state.gamma,
)

###########################################################################
# Initialize
###########################################################################
function Base.iterate(iter::ZEROFPR_iterable{R}) where {R}
    # initial point
    x = copy(iter.x0)
    g_xbar, x = objprox(iter.prob, x, iter.opts.bigeps)
    f_x, grad_f_x = objgrad(iter.prob, x)
    # scaling
    if iter.opts.grad_scaling
        scalobj = objgradscaling(
            grad_f_x,
            iter.opts.scaling_grad_target,
            iter.opts.scaling_grad_max,
            iter.opts.scaling_min_value,
        )
    else
        scalobj = R(1)
    end
    # forward-backward stepsize
    gamma = iter.gamma
    if gamma === nothing
        # compute lower bound to Lipschitz constant of x ↦ ∇f(x)
        xh = x .+ 0.01 .* (x .+ R(1))
        grad_f_xh = grad(iter.prob, xh)
        L = norm(scalobj .* (grad_f_xh - grad_f_x)) / R(sqrt(length(x)))
        L = max(1.0, L) # this avoids zero-division
        gamma = max(iter.opts.alpha / L, sqrt(iter.opts.gamma_min))
    end

    # initial forward-backward step
    y = x - gamma .* (scalobj .* grad_f_x)
    g_xbar, xbar = objprox(iter.prob, y, gamma * scalobj)

    # initial fixed-point residual and optimality
    res = x - xbar
    #optim = norm( res, Inf ) / gamma
    f_xbar, tmp = objgrad(iter.prob, xbar) # grad_f_xbar
    tmp .+= (res ./ gamma) .- grad_f_x
    optim = norm(tmp, Inf)

    # initial objective
    objec = f_xbar + g_xbar

    # forward-backward envelope
    f_xbar_upp = f_model(scalobj * f_x, scalobj .* grad_f_x, res, gamma)
    fbe_x = f_xbar_upp + scalobj * g_xbar

    state = ZEROFPR_state(
        x,
        f_x,
        grad_f_x,
        gamma,
        y,
        xbar,
        g_xbar,
        res,
        optim,
        iter.H,
        nothing,
        objec,
        scalobj,
        fbe_x,
        zero(x),
        zero(x),
        zero(x),
        zero(x),
        nothing,
        nothing,
    )

    return state, state
end

###########################################################################
# Step
###########################################################################
function Base.iterate(iter::ZEROFPR_iterable{R}, state::ZEROFPR_state{R,Tx}) where {R,Tx}

    # (scaled) quadratic model of `f` around `x` at `xbar`
    f_xbar_upp = f_model(state)

    # `f` at `xbar`
    f_xbar = objgrad!(iter.prob, state.xbar, state.grad_f_xbar)
    f_xbar *= state.scalobj

    # check and possibly backtrack stepsize `gamma`
    while iter.gamma === nothing || iter.opts.adaptive == true
        if state.gamma < iter.opts.gamma_min
            if iter.opts.warnings
                @warn "stepsize `gamma` too small ($(state.gamma))"
            end
            return nothing
        end
        tol = iter.opts.bigeps * (1 + abs(f_xbar))
        if f_xbar <= f_xbar_upp + tol
            break
        end
        state.gamma *= 0.5
        state.y .= state.x .- state.gamma .* (state.scalobj .* state.grad_f_x)
        state.g_xbar = objprox!(iter.prob, state.y, state.gamma * state.scalobj, state.xbar)
        state.res .= state.x .- state.xbar
        reset!(state.H)
        f_xbar_upp = f_model(state)
        f_xbar = objgrad!(iter.prob, state.xbar, state.grad_f_xbar)
        f_xbar *= state.scalobj
    end

    # (scaled) forward-backward envelope
    state.fbe_x = f_xbar_upp + state.scalobj * state.g_xbar

    # residual at `xbar`
    state.y .= state.xbar .- state.gamma .* (state.scalobj .* state.grad_f_xbar)
    g_xbarbar = objprox!(iter.prob, state.y, state.gamma * state.scalobj, state.xbarbar)
    state.res_xbar .= state.xbar .- state.xbarbar

    if state.xbar_old === nothing
        # store vectors for next update
        state.xbar_old = copy(state.xbar)
        state.res_xbar_old = copy(state.res_xbar)
    else
        # update variable metric
        update!(state.H, state.xbar - state.xbar_old, state.res_xbar - state.res_xbar_old)
        # store vectors for next update
        copyto!(state.xbar_old, state.xbar)
        copyto!(state.res_xbar_old, state.res_xbar)
    end

    # search direction
    mul!(state.d, state.H, -state.res_xbar)

    # linesearch over the forward-backward envelope
    tau = R(1)
    sigma = iter.opts.sigma / state.gamma
    tol = iter.opts.bigeps * (1 + abs(state.fbe_x))
    rhs = state.fbe_x - sigma * norm(state.res)^2 + tol

    while true
        state.x .= state.xbar_old .+ tau .* state.d
        state.f_x = objgrad!(iter.prob, state.x, state.grad_f_x)
        state.y .= state.x .- state.gamma .* (state.scalobj .* state.grad_f_x)
        state.g_xbar = objprox!(iter.prob, state.y, state.gamma * state.scalobj, state.xbar)
        state.res .= state.x .- state.xbar
        f_xbar_upp = f_model(state)
        state.fbe_x = f_xbar_upp + state.scalobj * state.g_xbar

        f_xbar = obj(iter.prob, state.xbar)
        f_xbar *= state.scalobj
        tol = iter.opts.bigeps * (1 + abs(f_xbar))
        is_gamma_ok = f_xbar <= f_xbar_upp + tol

        if (state.fbe_x <= rhs) && is_gamma_ok
            break

        elseif tau < iter.opts.tau_min
            if !is_gamma_ok
                state.gamma *= 0.5
                reset!(state.H)
            end
            objec_t = f_xbar + state.g_xbar
            tol = iter.opts.bigeps * (1 + abs(state.objec))
            if objec_t <= state.objec + tol
                # small linesearch stepsize : taking tau > 0, which gives decrease
                # in the objective function.
                break
            end
            # small linesearch stepsize : taking `tau = 0`, which gives decrease in
            # the FBE.
            tau = 0.0
            copyto!(state.x, state.xbar_old)
            state.f_x = objgrad!(iter.prob, state.x, state.grad_f_x)
            state.y .= state.x .- state.gamma .* (state.scalobj .* state.grad_f_x)
            state.g_xbar =
                objprox!(iter.prob, state.y, state.gamma * state.scalobj, state.xbar)
            state.res .= state.x .- state.xbar
            state.fbe_x = f_model(state) + state.scalobj * state.g_xbar
            break

        else
            tau *= 0.5
        end
    end

    state.tau = tau
    f_xbar, tmp = objgrad(iter.prob, state.xbar) # grad_f_xbar
    tmp .+= (state.res ./ state.gamma) .- state.grad_f_x
    state.optim = norm(tmp, Inf)
    state.objec = f_xbar + state.g_xbar
    return state, state
end

###########################################################################
# Solver
###########################################################################
mutable struct ZEROFPR{R<:Real}
    memory::Int
    max_iter::Int
    tol_optim::R
    verbose::Bool
    freq::Int
    tol_diverging::R
    opts::ZEROFPR_options

    function ZEROFPR{R}(;
        alpha::R = R(0.95),
        beta::R = R(0.5),
        adaptive::Bool = true,
        memory::Int = 8,
        max_iter::Int = 10000,
        tol_optim::R = R(1e-8),
        verbose::Bool = false,
        freq::Int = 10,
        tau_min::R = R(1e-8),
        gamma_min::R = R(1e-12),
        grad_scaling::Bool = true,
        scaling_grad_max::R = R(100),
        scaling_grad_target::Maybe{R} = nothing,
        scaling_min_value::R = R(1e-8),
        tol_diverging::R = R(1e20),
        bigeps::R = R(1e-12),
        warnings::Bool = true,
    ) where {R}
        @assert 0 < alpha < 1
        @assert 0 < beta < 1
        @assert 0 <= memory
        @assert 0 < max_iter
        @assert 0 < freq
        @assert 0 < bigeps < 1
        @assert 0 < tol_optim < 1 < tol_diverging
        @assert 0 < scaling_grad_max
        @assert scaling_grad_target === nothing || 0 < scaling_grad_target
        @assert 0 < scaling_min_value < 1
        @assert 0 < gamma_min < 1
        @assert 0 < tau_min < 1
        sigma = beta * 0.5 * (1.0 - alpha)
        opts = ZEROFPR_options(
            alpha,
            beta,
            sigma,
            adaptive,
            tau_min,
            gamma_min,
            bigeps,
            grad_scaling,
            scaling_grad_max,
            scaling_grad_target,
            scaling_min_value,
            warnings,
        )
        new(memory, max_iter, tol_optim, verbose, freq, tol_diverging, opts)
    end
end

###########################################################################
# Solver call
###########################################################################
function (solver::ZEROFPR{R})(
    prob::TP;
    x0::AbstractArray{R} = prob.meta.x0,
    L::Maybe{R} = nothing,
) where {R,TP<:AbstractOptiModel}

    tstart = time()

    stop_solved(state::ZEROFPR_state) = state.optim <= solver.tol_optim
    stop_unbound(state::ZEROFPR_state) = state.objec < -solver.tol_diverging
    is_illegal(x) = x === nothing ? false : (any(isnan.(x)) || any(isinf.(x)))
    stop_illegal(state::ZEROFPR_state) = is_illegal(state.x) || is_illegal(state.xbar)
    stop(state::ZEROFPR_state) =
        stop_solved(state) || stop_unbound(state) || stop_illegal(state)

    if solver.verbose
        @printf(
            "%6s | %8s | %8s  %9s | %9s  %9s\n",
            "iter",
            "objec",
            "fbe",
            "optim",
            "gamma",
            "tau"
        )
        disp((i, state)) = @printf(
            "%6d | %+.1e | %+.1e  %.3e | %.3e  %.3e \n",
            i,
            state.objec,
            state.fbe_x,
            state.optim,
            state.gamma,
            (state.tau === nothing ? 0.0 : state.tau)
        )
    end

    gamma = (L === nothing ? nothing : solver.opts.alpha / L)

    iter = ZEROFPR_iterable(prob, x0, gamma, LBFGS(x0, solver.memory), solver.opts)
    iter = halt(iter, stop)
    iter = take(iter, solver.max_iter)
    iter = enumerate(iter)
    if solver.verbose
        iter = sample(iter, solver.freq)
        iter = tee(iter, disp)
    end

    num_iters, state = loop(iter)

    status = if stop_solved(state)
        :first_order
    elseif stop_unbound(state)
        :unbounded
    elseif stop_illegal(state)
        :ieee_nan_inf
    elseif (num_iters >= solver.max_iter)
        :max_iter
    elseif (state.gamma <= solver.opts.gamma_min)
        :stalled
    else
        :unknown
    end

    return OptiOutput(
        status,
        x = state.xbar,
        objective = state.objec,
        optimality = state.optim,
        iterations = num_iters,
        time = time() - tstart,
        solver_name = "ZEROFPR",
        solver = Dict(:gx => state.g_xbar),
    )
end

###########################################################################
# Outer constructors
###########################################################################
"""
    ZEROFPR([gamma, adaptive, memory, max_iter, tol_optim, verbose, freq, ...])

Instantiate the ZeroFPR algorithm [TSP18] for solving optimization problems
of the form

    minimize f(x) + g(x),

where `f` is smooth and `g` is proximable. If `solver = ZEROFPR(args...)`, then the
above problem is solved with

    solver( prob, [x0, L] )

Optional keyword arguments:

* `gamma::Real` (default: `nothing`), the stepsize to use; defaults to `alpha/L` if not set (but `L` is).
* `adaptive::Bool` (default: `true`), if true, forces the method stepsize to be adaptively adjusted even
    if `L` is provided (this behaviour is always enforced if `L` is not provided).
* `memory::Integer` (default: `5`), memory parameter for L-BFGS.
* `max_iter::Integer` (default: `1000`), maximum number of iterations to perform.
* `tol_optim::Real` (default: `1e-8`), absolute tolerance on the fixed-point residual.
* `verbose::Bool` (default: `true`), whether or not to print information during the iterations.
* `freq::Integer` (default: `10`), frequency of verbosity.
* `alpha::Real` (default: `0.95`), stepsize to inverse-Lipschitz-constant ratio; should be in (0, 1).
* `beta::Real` (default: `0.5`), sufficient decrease parameter; should be in (0, 1).
* `grad_scaling::Bool`
* `scaling_grad_max::Real`
* `scaling_grad_target::Maybe{Real}`
* `scaling_min_value::Real`

If `gamma` is not specified at construction time, the following keyword
argument can be used to set the stepsize parameter:

* `L::Real` (default: `nothing`), the Lipschitz constant of the gradient of x ↦ f(Ax).

References:

[TSP18] Themelis, Stella, Patrinos, ``Forward-backward envelope for the sum of two
        nonconvex functions: Further properties and nonmonotone line-search algorithms``,
        SIAM Journal on Optimization, vol. 28, no. 3, pp. 2274–2303 (2018).
        DOI: 10.1137/16M1080240
"""
ZEROFPR(::Type{R}; kwargs...) where {R} = ZEROFPR{R}(; kwargs...)
ZEROFPR(; kwargs...) = ZEROFPR(Float64; kwargs...)


"""
    QuarticNonconvex{T}

The scalar conservation law
``
  \\partial_t u + \\partial_x ( u(t,x)^4 - 10 u(t,x)^2 + 3 u(t,x) ) = 0
``
in one space dimensions using `T` as scalar type.
"""
struct QuarticNonconvex{T} <: ScalarBalanceLaw{T,1} end

function QuarticNonconvex(T=Float64)
  QuarticNonconvex{T}()
end

function Base.show(io::IO, model::QuarticNonconvex{T}) where {T}
  print(io, "Scalar conservation law {T=", T, "} with flux f(u) = u^4 - 10 u^2 + 3 u")
end

"""
    flux(u, model::QuarticNonconvex)

Compute the flux of `u` for `model`.
"""
@inline flux(u, model::QuarticNonconvex) = u^2 * (u^2 - 10) + 3u

"""
    speed(u::Real, model::QuarticNonconvex)

Compute the speed f'(`u`) for `model`.
"""
@inline speed(u, model::QuarticNonconvex) = 4*u^3 - 20*u + 3

function min_max_speed(ul::T, ur::T, model::QuarticNonconvex{T}) where T<:Real
  min_speed = zero(T)
  max_speed = zero(T)

  if ur > 1.2909944487358056284 && -2.5819888974716112568 < ul < 1.2909944487358056284
    min_speed = T(-14.213259316477408379)
  elseif (ur > 1.2909944487358056284 && (ul >= 1.2909944487358056284 || ul <= -2.5819888974716112568)) || (ur <= 1.2909944487358056284 && 2 * ul + ur + sqrt(20 - 3 * max(ur, 1.2909944487358056284)^2) <= 0) || 1.2909944487358056284 + ur <= 0
    min_speed = 3 + 4 * ul * (-5 + ul^2)
  else
    min_speed = 3 + 4 * ur * (-5 + ur^2)
  end

  if (1.2909944487358056284 + ul <= 0 && 1.2909944487358056284 + ur > 0 && ur < 1.2909944487358056284) || (1.2909944487358056284 + ul < 0 && 1.2909944487358056284 <= ur < 2.5819888974716112568)
    max_speed = T(20.213259316477408379)
  elseif (1.2909944487358056284 + ul == 0 && ur == 2.5819888974716112568) || (1.2909944487358056284 + ul >= 0 && (ur ≈ 1.2909944487358056284 || (1.2909944487358056284 < ur < 2.5819888974716112568 && sqrt(20 - 3 * min(ur, 2.5819888974716112568)^2) >= 2 * ul + ur))) || (ur < 1.2909944487358056284 && 1.2909944487358056284 + ul > 0)
    max_speed = 3 + 4 * ul * (-5 + ul^2)
  else
    max_speed = 3 + 4 * ur * (-5 + ur^2)
  end

  min_speed, max_speed
end

function max_abs_speed(ul::T, ur::T, model::QuarticNonconvex{T}) where T<:Real
  min_speed, max_speed = min_max_speed(ul, ur, model)
  max(abs(min_speed), abs(max_speed))
end

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

"""
    (::EnergyConservativeFlux)(uₗ::T, uᵣ::T, model::QuarticNonconvex{T}) where {T}

Compute the energy (L₂ entropy) conservative flux between `uₗ` and `uᵣ` for `model`.
"""
function (flux::EnergyConservativeFlux)(ul::T, ur::T, model::QuarticNonconvex{T}) where T<:Real
  (ul^4 + ur*ul^3 + ur^2*ul^2 + ur^3*ul + ur^4) / 5 - 10 * (ul^2 + ul*ur + ur^2) / 3 + 3 * (ul + ur) / 2
end


# Archimedean copulas
"""
  getV0(θ::Union{Float64, Int}, v::Vector{Float64}, copula::String)

Returns Vector{Float} or Vector{Int} of realisations of axiliary variable V0
used to ganarate data from 1d archimedean copula with parameter Θ, given v:
realisations of 1d variable uniformly distributed on [0,1]

```jldoctest

julia> getV0(2., [0.2, 0.4, 0.6, 0.8], "clayton")
4-element Array{Float64,1}:
 0.0641848
 0.274996
 0.708326
 1.64237
```
"""
function getV0(θ::Union{Float64, Int}, v::Vector{Float64}, copula::String)
  if copula == "clayton"
    return quantile.(Gamma(1/θ, 1), v)
  elseif copula == "amh"
    return 1 .+quantile.(Geometric(1-θ), v)
  elseif copula == "frank"
    return logseriesquantile(1-exp(-θ), v)
  elseif copula == "gumbel"
    return levygen(θ, v)
  end
  throw(AssertionError("$(copula) not supported"))
end

"""
  phi(u::Matrix{Float64}, θ::Union{Float64, Int}, copula::String)

Given a matrix t realisations of n variate data ℜᵗⁿ ∋ u = -log(rand(t,n))./V0
returns it transformed through an inverse generator of archimedean copula. Output
is distributed uniformly on [0,1]ⁿ

```jldoctest

julia> julia> phi([0.2 0.6; 0.4 0.8], 2., "clayton")
2×2 Array{Float64,2}:
 0.845154  0.6742
 0.745356  0.620174
```
"""
function phi(u::Matrix{Float64}, θ::Union{Float64, Int}, copula::String)
  if copula == "clayton"
    return (1 .+ u).^(-1/θ)
  elseif copula == "amh"
    return (1-θ) ./(exp.(u) .-θ)
  elseif copula == "frank"
    return -log.(1 .+exp.(-u) .*(exp(-θ)-1)) ./θ
  elseif copula == "gumbel"
    return exp.(-u.^(1/θ))
  end
  throw(AssertionError("$(copula) not supported"))
end

"""
  copulagen(copula::String, r::Matrix{Float}, θ::Union{Float64, Int})

Auxiliary function used to generate data from archimedean copula (clayton, gumbel, frank or amh)
parametrised by a single parameter θ given a matrix of independent [0,1] distributerd
random vectors.

```jldoctest
  julia> copulagen("clayton", [0.2 0.6 0.9; 0.4 0.5 0.8], 2.)
  2×2 Array{Float64,2}:
   0.675778  0.851993
   0.687482  0.736394
```

"""
function copulagen(copula::String, r::Matrix{T}, θ::Union{Float64, Int}) where T <:AbstractFloat
  u = -log.(r[:,1:end-1])./getV0(θ, r[:,end], copula)
  phi(u, θ, copula)
end

"""

  archcopulagen(t::Int, n::Int, θ::Union{Float64, Int}, copula::String; rev::Bool = false)

Returns: t x n Matrix{Float}, t realisations of n-variate data generated from Archimedean
one parameter copula.

Following copula families are supported: clayton, frank, gumbel and amh --
Ali-Mikhail-Haq.

If rev == true, reverse the copula output i.e. u → 1-u (we call it reversed copula).
It cor == pearson, kendall, uses correlation coeficient as a parameter

```jldoctest
julia> Random.seed!(43);

julia> archcopulagen(10, 2, 1, "clayton")
10×2 Array{Float64,2}:
 0.770331  0.932834
 0.472847  0.0806845
 0.970749  0.653029
 0.622159  0.0518025
 0.402461  0.228549
 0.946375  0.842883
 0.809076  0.129038
 0.747983  0.433829
 0.374341  0.437269
 0.973066  0.910103

 ```
"""
function archcopulagen(t::Int, n::Int, θ::Union{Float64, Int}, copula::String;
                                                              rev::Bool = false,
                                                              cor::String = "")
  cor in ["Spearman", "Kendall", ""] || throw(AssertionError("$(cor) correlation not supported"))
  copula in ["clayton", "amh", "frank", "gumbel"] || throw(AssertionError("$(copula) copula is not supported"))
  if (n == 2)*((copula != "gumbel")*(θ < 0) | (copula == "amh")*(θ in [0,1]))
    return chaincopulagen(t, [θ], copula; rev=rev, cor=cor)
  elseif cor == "Spearman"
    θ = useρ(θ , copula)
  elseif cor == "Kendall"
    θ = useτ(θ , copula)
  else
    testθ(θ, copula)
  end
  u = copulagen(copula, rand(t,n+1), θ)
  rev ? 1 .-u : u
end

"""
  function testθ(θ::Union{Float64, Int}, copula::String)

Tests the parameter θ value for archimedean copula, returns void
"""
function testθ(θ::Union{Float64, Int}, copula::String)
  if copula == "gumbel"
    θ >= 1 || throw(DomainError("generaton not supported for θ < 1"))
  elseif copula == "amh"
    1 > θ > 0 || throw(DomainError("amh multiv. copula supported only for 0 < θ < 1"))
  else
    θ > 0 || throw(DomainError("generaton not supported for θ ≤ 0"))
  end
  Nothing
end

"""
  useρ(ρ::Float64, copula::String)

Tests the available pearson correlation for archimedean copula, returns Float,
corresponding copula parameter θ.
```jldoctest

julia> useρ(0.75, "gumbel")
2.294053859606698
```

"""
function useρ(ρ::Float64, copula::String)
  0 < ρ < 1 || throw(DomainError("correlation coeficiant must fulfill 0 < ρ < 1"))
  if copula == "amh"
    0 < ρ < 0.5 || throw(DomainError("correlation coeficiant must fulfill 0 < ρ < 0.5"))
  end
  ρ2θ(ρ, copula)
end

"""
  useτ(ρ::Float64, copula::String)

Tests the available kendall's correlation for archimedean copula, returns Float,
corresponding copula parameter θ.

```jldoctest

julia> useτ(0.5, "clayton")
2.0
```
"""
function useτ(τ::Float64, copula::String)
  0 < τ < 1 || throw(DomainError("correlation coeficiant must fulfill 0 < τ < 1"))
  if copula == "amh"
    0 < τ < 1/3 || throw(DomainError("correlation coeficiant must fulfill 0 < τ < 1/3"))
  end
  τ2θ(τ, copula)
end


using Routing
using Test

include("debugging.jl")

test_start_time = time()

include("espprc_test.jl")
include("vrptw_test.jl")

@testset "Solomon Instance Test" begin
    solomon_dataset = Dict(
        "C101_025" => 191.30,
        "C102_025" => 190.30,
        "C201_025" => 214.70,
        "RC105_025" => 411.30,
        "R101_025" => 617.10,
        "R102_025" => 547.10,
    )

    for (name, val) in solomon_dataset
        @info("Solving the Solomon Instance $name")
        solomon = load_solomon(name)
        vrptw = generate_solomon_vrptw_instance(solomon)

        @testset "$name" begin
            println("-- Solomon Instance $name --")
            @time routes, objective_value = solveVRP(vrptw, pricing_method="pulse");
            @test isapprox(objective_value, val, atol=1e-7)
        end
    end
end

total_test_time = time() - test_start_time
@show total_test_time


function check_cutval(A,svec,tvec,S)
    """
    Check the cut value associated with the set of nodes in graph A.
    Source and sink nodes are not explicitly included, but terminal
    edges weights are given in svec and tvec.
    """
    cut = 0
    n = size(A,1)
    Sbar = setdiff(1:n,S)   # complement set

    # terminal edge weight penalties
    cut += sum(tvec[S])
    cut += sum(svec[Sbar])
    AS = A[S,Sbar]
    cut += sum(AS)

    return cut
end

function clique_expansion_penalties(k)
    """
    Symmetric clique splitting penalty for hyperedge of size k.
    Only gives first r = floor(k/2) penalties
    """
    @assert(k > 1)
    r = floor(Int64,(k)/2)
    I = collect(0:k)
    w = I .* (k .- I)
    w = w[1:r+1]
    return w
end

function load_superpixel_cliques(dataset,lambda,numsuperpix=500)
    """
    Load superpixel regions for a certain image dataset from experiements of
    Jegalka et al.

    Save the superpixel regions in hyperedge list format, with clique
    expansion splitting penalties, weighted by lambda
    (use lambda values as given previously by Jegelka et al.)
    """

    # Load superpixel region information
    M = matread("../data/data_reflections/$(dataset)_p$numsuperpix.mat")
    sup_img = M["sup_img"]

    col = size(sup_img,2)                   # number of columns
    row = size(sup_img,1)                   # number of rows
    node = (i,j)->i + (j-1)*row             # go from (i,j) node ID to linear node ID
    region = length(unique(sup_img))        # number of superpixels

    # Save superpixels in alternate format
    Edges = Vector{Vector{Int64}}()
    for e = 1:region
        # each superpixel defines a hyperedge
        push!(Edges,Vector{Int64}())
    end

    EdgeLengths = zeros(Int64,region)
    for i = 1:row
        for j = 1:col
            v = node(i,j)    # linear index of node at gridpoint (i,j)

            # the superpixel, i.e., hyperedge, it's in
            e = round(Int64,sup_img[i,j])
            push!(Edges[e],v)
            EdgeLengths[e] += 1
        end
    end

    # Discard cliques with less than 2 nodes, they do nothing
    keep = findall(x->x>1,EdgeLengths)
    Edges = Edges[keep]
    EdgesW = Vector{Vector{Float64}}()
    for ei = 1:length(Edges)
        # We'll use clique expansion splitting penalties,
        # scaled with previously determined weights of Jegelka et al.
        e = Edges[ei]
        push!(EdgesW,lambda*clique_expansion_penalties(length(e)))
    end

    return Edges, EdgesW
end

function eval_fS_image(Edges,EdgesW,svec,tvec,A,S)
    """
    For a set S, evaluate the objective from a decomposable
    submodular function involving terminal edges as defined
    in svec and tvec, pairwise non-terminal edges
    as defined in A, and higher-order terms from Edges
    and EdgesW.
    """
    n = size(A,1)
    eS = zeros(n)
    eS[S] = ones(length(S))
    Sbar = setdiff(1:n,S)   # complement set

    # terminal edge weight penalties
    cutval = sum(tvec[S])
    cutval += sum(svec[Sbar])

    # interior edge penalties
    AS = A[S,Sbar]
    cutval += sum(AS)

    # higher-order cut penalties
    for i = 1:length(Edges)
        e = Edges[i]
        w = EdgesW[i]
        t = round(Int64,sum(eS[e]))

        # do this because it's a symmetric splitting function,
        # so w is only half the size
        t = min(t,length(e)-t)
        cutval += w[t+1]
    end

    return cutval

end

function sweep_cut_f(Edges,EdgesW,svec,tvec,A,y)
    if length(y) > 10
        println("This code is inefficient. Not recommended for use for vectors this big.")
        return
    end
     Order = sortperm(y,rev = true)
     n = length(y)
     @assert(n == size(A,1))
     Sbest = [Order[1]]
     fbest = eval_fS_image(Edges,EdgesW,svec,tvec,A,Sbest)
     for i = 2:n
         S = Order[1:i]
         f = eval_fS_image(Edges,EdgesW,svec,tvec,A,S)
         if f < fbest
             fbest = f
             Sbest = S
         end
     end

     return Sbest, fbest
 end

function Objective_from_Primals(Xmat,sweepi,Edges,EdgesW,svec,tvec,An)
    """
    Given a set of primal variables from a continuous DSFM method,
    stored as columns of Xmat, find the objective score resulting
    from the best level set S_lam = {i : x_i >= lam}.

    The value of sweepi tells us the best level set as
    determined by the Matlab/C++ code that Xmat comes from.
    """
    sourceside = sum(svec)
    t = size(Xmat,2)
    @assert(t == length(sweepi))
    Fvals = zeros(t)
    for i = 1:t
        x = Xmat[:,i]
        Order = sortperm(x,rev=true)
        wherei = findfirst(x->x==sweepi[i]+1,Order)
        S = Order[1:wherei]
        Fvals[i] = eval_fS_image(Edges,EdgesW,svec,tvec,An,S)
    end
    return Fvals
end


"""
```
blanchard_kahn(m::RiskAdjustedLinearization; deterministic::Bool = false, verbose::Symbol = :high)
```

checks the Blanchard-Kahn conditions for whether a first-order perturbation is saddle-path stable or not.

If `verbose` is `:low` or `:high`, a print statement will be shown if the Blanchard-Kahn conditions are satisfied.
"""
function blanchard_kahn(m::RiskAdjustedLinearization; deterministic::Bool = false, verbose::Symbol = :high)

    li = linearized_system(m)

    Nz = m.Nz
    Ny = m.Ny
    N_zy = m.Nz + m.Ny
    ztype = eltype(m.z)
    AA = Matrix{ztype}(undef, N_zy, N_zy)
    BB = similar(AA)

    # Populate AA
    AA[1:Ny, 1:Nz]                 = li[:Γ₅]
    AA[1:Ny, (Nz + 1):end]         = li[:Γ₆]
    AA[(Ny + 1):end, 1:Nz]         = Matrix{ztype}(I, m.Nz, m.Nz) # faster but makes allocations, unlike Diagonal(Ones{ztype}(Nz))
    AA[(Ny + 1):end, (Nz + 1):end] = Zeros{ztype}(m.Nz, m.Ny)

    # Populate BB
    BB[1:Ny, 1:Nz]                 = deterministic ? -li[:Γ₃] : -(li[:Γ₃] + li[:JV])
    BB[1:Ny, (Nz + 1):end]         = -li[:Γ₄]
    BB[(Ny + 1):end, 1:Nz]         =  li[:Γ₁]
    BB[(Ny + 1):end, (Nz + 1):end] =  li[:Γ₂]

    if count(abs.(eigen(AA, BB).values) .> 1) != m.Nz
        if deterministic
            throw(BlanchardKahnError("First-order perturbation around deterministic steady state is not saddle-path stable"))
        else
            throw(BlanchardKahnError("First-order perturbation around stochastic steady state is not saddle-path stable"))
        end
    else
        if verbose in [:low, :high]
            if deterministic
                println("Blanchard-Kahn conditions for a unique locally bounded deterministic " *
                        "steady-state perturbation are satisfied")
            else
                println("Blanchard-Kahn conditions for a unique locally bounded stochastic " *
                        "steady-state perturbation are satisfied")
            end
        end

        return true
    end
end

mutable struct BlanchardKahnError <: Exception
    msg::String
end
BlanchardKahnError() = BlanchardKahnError("First-order perturbation is not saddle-path stable")
Base.showerror(io::IO, ex::BlanchardKahnError) = print(io, ex.msg)


module Categories

export Category, Object, Morphism, Composition, Ob, Hom, category, id, label, domain, codomain, compose

struct Object{O}
    object::O
end

obtype(::Object{O}) where {O} = O

abstract type AbstractMorphism{O,M} end

obtype(::AbstractMorphism{O}) where {O} = O

homtype(::AbstractMorphism{O,M}) where {O,M} = M

struct Id{O,M} <: AbstractMorphism{O,M} end
(id::Id{O})(x::O) where {O} = x

struct Morphism{O,M} <: AbstractMorphism{O,M}
    morphism::M
    domain::Object{O}
    codomain::Object{O}
end

morphism(m::Morphism) = m.morphism

domain(m::Morphism) = m.domain

codomain(m::Morphism) = m.codomain

struct Composition{O,M,N} <: AbstractMorphism{O,M}
    morphisms::NTuple{N,Morphism{O,M}}
end

domain(c::Composition) = first(c.morphisms).domain

codomain(c::Composition) = last(c.morphisms).codomain

struct Hom{O,M}
    domain::Object{O}
    codomain::Object{O}
end

domain(hom::Hom) = hom.domain

codomain(hom::Hom) = hom.codomain

struct Category{O,M}
    ob::Type{Object{O}}
    hom::Type{Morphism{O,M}}
    id::Id{O,M}

    Category(::Type{O},::Type{M}) where {O,M} = new{O,M}(Object{O},Morphism{O,M},Id{O,M}())
end

Ob(C::Category) = C.ob

Hom(C::Category)= C.hom

Hom(::Category{O,M},X::Object{O},Y::Object{O}) where {O,M} = Hom{O,M}(X,Y)

id(C::Category) = C.id

Base.in(::Object{O1}, ::Type{Object{O2}}) where {O1,O2} = O1 <: O2

Base.in(::AbstractMorphism{O1,M1}, ::Type{Morphism{O2,M2}}) where {O1,O2,M1,M2} = (O1 <: O2) && (M1 <: M2)

function Base.in(m::AbstractMorphism{O1,M1}, hom::Hom{O2,M2}) where {O1,O2,M1,M2}
    (O1 <: O2) &&
    (M1 <: M2) &&
    (domain(m) === domain(hom)) &&
    (codomain(m) === codomain(hom))
end

function compose(m1::Morphism{O,M},m2::Morphism{O,M}) where {O,M}
    m1.codomain === m2.domain || error("Domain mismatch")
    return Composition((m1,m2))
end

function compose(m::Morphism{O,M},c::Composition{O,M}) where {O,M}
    m.codomain === domain(c) || error("Domain mismatch")
    Composition((m,c.morphisms...))
end

function compose(c::Composition{O,M},m::Morphism{O,M}) where {O,M}
    codomain(c) === m.domain || error("Domain mismatch")
    Composition((c.morphisms...,m))
end

compose(m::AbstractMorphism{O,M},::Id{O,M}) where {O,M} = m

compose(::Id{O,M},m::AbstractMorphism{O,M}) where {O,M} = m

Base.:∘(m1::AbstractMorphism,m2::AbstractMorphism) = compose(m1,m2)

function (comp::Composition{O})(x::O) where {O}
    domain(comp).object === x || error("Domain mismatch")
    for mor in comp.morphisms
        x = mor(x)
    end
    return x
end

Base.show(io::IO, obj::Object) = print(io, obj.object)

Base.show(io::IO, mor::Morphism) = print(io, mor.morphism, ": ", mor.domain," → ", mor.codomain)

function Base.show(io::IO, comp::Composition)
    print(
        io,
        join(morphism.(comp.morphisms)," ∘ "),
        ": ",
        domain(first(comp.morphisms)),
        " → ",
        codomain(last(comp.morphisms))
    )
end

end # module


push!(LOAD_PATH, "../models")
using SuperResModels
using SparseInverseProblems

@everywhere begin
    include("../ufus_parameters.jl")
    parameters = SuperResModels.Conv2dParameters(x_max, x_max, filter, filter_dx, filter_dy, sigma, sigma, sigma, sigma)
    model_static = SuperResModels.Conv2d(parameters)
    dynamic_parameters = SuperResModels.DynamicConv2dParameters(K, tau, v_max, 20)
    model_dynamic = SuperResModels.DynamicConv2d(model_static, dynamic_parameters)
    n_x = model_static.n_x
    n_y = model_static.n_y
end

# Get time of script start
now_str = string(now())
now_str = replace(now_str, ":", "-")
now_str = replace(now_str, ".", "-")

# L2 norm single particle
thetas = reshape([x_max/2; x_max/2], 2, 1)
weights = [1.0]
single_particle_norm = norm(phi(model_static, thetas, weights), lp_norm)
@everywhere single_particle_norm = $single_particle_norm
println("norm = ", single_particle_norm)

# Test dynamic
println("Testing dynamic...")
particles_m = [x_max/4]
#particles_p = [3*x_max/4]
thetas = hcat([(x_max/2 - dx) * ones(1, length(particles_m)); particles_m'; 0.0; -v_max/2])
              #[(x_max/2 + dx) * ones(1, length(particles_p)); particles_p'; 0.0; v_max/2])
weights = ones(1)
target = phi(model_dynamic, thetas, weights)
function callback(old_thetas, thetas, weights, output, old_obj_val)
    #evalute current OV
    new_obj_val,t = SparseInverseProblems.loss(SparseInverseProblems.LSLoss(), output - target)
    #println("gap = $(old_obj_val - new_obj_val)")
    if old_obj_val - new_obj_val < 1E-4
        return true
    end
    return false
end
(thetas_est,weights_est) = SparseInverseProblems.ADCG(model_dynamic, SparseInverseProblems.LSLoss(), target, 2.0, callback=callback, max_iters=200)
println("theta = ", thetas, ", weights = ", weights)
println("theta_est = ", thetas_est, ", weights_est = ", weights_est)


using Test
using DelimitedFiles, LinearAlgebra, Statistics #, MLDatasets

using StableRNGs
#rng = StableRNG(123)
using BetaML

TESTRNG = FIXEDRNG # This could change...
#TESTRNG = StableRNG(123)

println("*** Testing Neural Network...")

# ==================================
# New test
# ==================================
println("Testing Learnable structure...")

L1a = Learnable(([1.0 2; 3 4.0], [1.1,2.2], Float64[], [1.1,2.2], [1.1]))
L1b = Learnable(([1.0 2; 3 4.0], [1.1,2.2], Float64[], [1.1,2.2], [1.1]))
L1c = Learnable(([1.0 2; 3 4.0], [1.1,2.2], Float64[], [1.1,2.2], [1.1]))
foo = (((((sum(L1a,L1b,L1c) - (L1a + L1b + L1c)) + 10) - 4) * 2 ) / 3) *
      (((((sum(L1a,L1b,L1c) - (L1a + L1b + L1c)) + 10) - 4) * 2 ) / 3)

@test foo.data[1][2] == 16.0 && foo.data[5][1] == 16.00
@test (L1a -1).data[1] ≈ (-1 * (1 - L1a)).data[1]
@test (2 / (L1a / 2)).data[2] ≈ (4/L1a).data[2]
@test sqrt(L1a).data[1][2,2] == 2.0

# ==================================
# New test
# ==================================
println("Testing basic NN behaviour...")
x = [0.1,1]
y = [1,0]
l1   = DenseNoBiasLayer(2,2,w=[2 1;1 1],f=identity,rng=copy(TESTRNG))
l2   = VectorFunctionLayer(2,f=softmax)
mynn = buildNetwork([l1,l2],squaredCost,name="Simple Multinomial logistic regression")
o1 = forward(l1,x)
@test o1 == [1.2,1.1]
#@code_warntype forward(l1,x)
o2 = forward(l2,o1)
@test o2 ≈ [0.5249791874789399, 0.47502081252106] ≈ softmax([1.2,1.1])
orig = predict(mynn,x')[1,:]
@test orig == o2
#@code_warntype Nn.predict(mynn,x')[1,:]
ϵ = squaredCost(o2,y)
#@code_warntype  squaredCost(o2,y)
lossOrig = loss(mynn,x',y')
@test ϵ == lossOrig
#@code_warntype  loss(mynn,x',y')
dϵ_do2 = dSquaredCost(o2,y)
@test dϵ_do2 == [-0.4750208125210601,0.47502081252106]
#@code_warntype dSquaredCost(o2,y)
dϵ_do1 = backward(l2,o1,dϵ_do2)
@test dϵ_do1 ≈ [-0.23691761847142412, 0.23691761847142412]
#@code_warntype backward(l2,o1,dϵ_do2)
dϵ_dX = backward(l1,x,dϵ_do1)
@test dϵ_dX ≈ [-0.23691761847142412, 0.0]
#@code_warntype backward(l1,x,dϵ_do1)
l1w = getParams(l1)
#@code_warntype
l2w = getParams(l2)
#@code_warntype
w = getParams(mynn)
#@code_warntype getParams(mynn)
#typeof(w)2-element Vector{Float64}:
origW = deepcopy(w)
l2dw = getGradient(l2,o1,dϵ_do2)
@test length(l2dw.data) == 0
#@code_warntype getGradient(l2,o1,dϵ_do2)
l1dw = getGradient(l1,x,dϵ_do1)
@test l1dw.data[1] ≈ [-0.023691761847142414 -0.23691761847142412; 0.023691761847142414 0.23691761847142412]
#@code_warntype
dw = getGradient(mynn,x,y)
#@code_warntype getGradient(mynn,x,y)
y_deltax1 = predict(mynn,[x[1]+0.001 x[2]])[1,:]
#@code_warntype Nn.predict(mynn,[x[1]+0.001 x[2]])[1,:]
lossDeltax1 = loss(mynn,[x[1]+0.001 x[2]],y')
#@code_warntype
deltaloss = dot(dϵ_dX,[0.001,0])
#@code_warntype
@test isapprox(lossDeltax1-lossOrig,deltaloss,atol=0.0000001)
l1wNew = l1w
l1wNew.data[1][1,1] += 0.001
setParams!(l1,l1wNew)
lossDelPar = loss(mynn,x',y')
#@code_warntype
deltaLossPar = 0.001*l1dw.data[1][1,1]
lossDelPar - lossOrig
@test isapprox(lossDelPar - lossOrig,deltaLossPar,atol=0.00000001)
η = 0.01
#w = gradientDescentSingleUpdate(w,dw,η)
#w = w - dw * η
w = w - dw * η
#@code_warntype gradSub.(w, gradMul.(dw,η))
#@code_warntype
setParams!(mynn,w)
loss2 = loss(mynn,x',y')
#@code_warntype
@test loss2 < lossOrig
for i in 1:10000
    local w  = getParams(mynn)
    local dw = getGradient(mynn,x,y)
    w  = w - dw * η
    setParams!(mynn,w)
end
lossFinal = loss(mynn,x',y')
@test predict(mynn,x')[1,1]>0.96
setParams!(mynn,origW)
train!(mynn,x',y',epochs=10000,batchSize=1,sequential=true,verbosity=NONE,optAlg=SGD(η=t->η,λ=1),rng=copy(TESTRNG))
#@code_warntype train!(mynn,x',y',epochs=10000,batchSize=1,sequential=true,verbosity=NONE,optAlg=SGD(η=t->η,λ=1))
lossTraining = loss(mynn,x',y')
#@code_warntype
@test isapprox(lossFinal,lossTraining,atol=0.00001)

li   = DenseLayer(2,2,w=[2 1;1 1],f=identity,rng=copy(TESTRNG))
@test getNParams(li) == 6

# ==================================
# NEW Test
println("Testing regression if it just works with manual derivatives...")

xtrain = [0.1 0.2; 0.3 0.5; 0.4 0.1; 0.5 0.4; 0.7 0.9; 0.2 0.1]
ytrain = [0.3; 0.8; 0.5; 0.9; 1.6; 0.3]
xtest = [0.5 0.6; 0.14 0.2; 0.3 0.7; 2.0 4.0]
ytest = [1.1; 0.36; 1.0; 6.0]
l1 = DenseLayer(2,3,w=[1 1; 1 1; 1 1], wb=[0 0 0], f=tanh, df=dtanh,rng=copy(TESTRNG))
l2 = DenseNoBiasLayer(3,2, w=[1 1 1; 1 1 1], f=relu, df=drelu,rng=copy(TESTRNG))
l3 = DenseLayer(2,1, w=[1 1], wb=[0], f=identity,df=didentity,rng=copy(TESTRNG))
mynn = buildNetwork([l1,l2,l3],squaredCost,name="Feed-forward Neural Network Model 1",dcf=dSquaredCost)
train!(mynn,xtrain,ytrain,batchSize=1,sequential=true,epochs=100,verbosity=NONE,optAlg=SGD(η=t -> 1/(1+t),λ=1),rng=copy(TESTRNG))
#@benchmark train!(mynn,xtrain,ytrain,batchSize=1,sequential=true,epochs=100,verbosity=NONE,optAlg=SGD(η=t -> 1/(1+t),λ=1))
avgLoss = loss(mynn,xtest,ytest)
@test  avgLoss ≈ 1.599729991966362
expectedOutput = [0.7360644412052633, 0.7360644412052633, 0.7360644412052633, 2.47093434438514]
predicted = dropdims(predict(mynn,xtest),dims=2)
@test any(isapprox(expectedOutput,predicted))

# With the ADAM optimizer...
l1 = DenseLayer(2,3,w=[1 1; 1 1; 1 1], wb=[0 0 0], f=tanh, df=dtanh,rng=copy(TESTRNG))
l2 = DenseNoBiasLayer(3,2, w=[1 1 1; 1 1 1], f=relu, df=drelu,rng=copy(TESTRNG))
l3 = DenseLayer(2,1, w=[1 1], wb=[0], f=identity,df=didentity,rng=copy(TESTRNG))
mynn = buildNetwork([l1,l2,l3],squaredCost,name="Feed-forward Neural Network with ADAM",dcf=dSquaredCost)
train!(mynn,xtrain,ytrain,batchSize=1,sequential=true,epochs=100,verbosity=NONE,optAlg=ADAM(),rng=copy(TESTRNG))
avgLoss = loss(mynn,xtest,ytest)
@test  avgLoss ≈ 0.9497779759064725
expectedOutput = [1.7020525792404175, -0.1074729043392682, 1.4998367847079956, 3.3985794704732717]
predicted = dropdims(predict(mynn,xtest),dims=2)
@test any(isapprox(expectedOutput,predicted))

# ==================================
# NEW TEST
# ==================================
println("Testing using AD...")

ϵtrain = [1.023,1.08,0.961,0.919,0.933,0.993,1.011,0.923,1.084,1.037,1.012]
ϵtest  = [1.056,0.902,0.998,0.977]
xtrain = [0.1 0.2; 0.3 0.5; 0.4 0.1; 0.5 0.4; 0.7 0.9; 0.2 0.1; 0.4 0.2; 0.3 0.3; 0.6 0.9; 0.3 0.4; 0.9 0.8]
ytrain = [(0.1*x[1]+0.2*x[2]+0.3)*ϵtrain[i] for (i,x) in enumerate(eachrow(xtrain))]
xtest  = [0.5 0.6; 0.14 0.2; 0.3 0.7; 20.0 40.0;]
ytest  = [(0.1*x[1]+0.2*x[2]+0.3)*ϵtest[i] for (i,x) in enumerate(eachrow(xtest))]

l1   = DenseLayer(2,3,w=ones(3,2), wb=zeros(3),rng=copy(TESTRNG))
l2   = DenseLayer(3,1, w=ones(1,3), wb=zeros(1),rng=copy(TESTRNG))
mynn = buildNetwork([l1,l2],squaredCost,name="Feed-forward Neural Network Model 1")
train!(mynn,xtrain,ytrain,epochs=1000,sequential=true,batchSize=1,verbosity=NONE,optAlg=SGD(η=t->0.01,λ=1),rng=copy(TESTRNG))
avgLoss = loss(mynn,xtest,ytest)
@test  avgLoss ≈ 0.0032018998005211886
ŷtestExpected = [0.4676699631752518,0.3448383593117405,0.4500863419692639,9.908883999376018]
ŷtrain = dropdims(predict(mynn,xtrain),dims=2)
ŷtest = dropdims(predict(mynn,xtest),dims=2)
@test any(isapprox(ŷtest,ŷtestExpected))
mreTrain = meanRelError(ŷtrain,ytrain)
@test mreTrain <= 0.06
mreTest  = meanRelError(ŷtest,ytest)
@test mreTest <= 0.05

#predicted = dropdims(Nn.predict(mynn,xtrain),dims=2)
#ytrain

# ==================================
# NEW TEST
# ==================================
println("Going through Multinomial logistic regression (using softmax)...")
#=
using RDatasets
using Random
using DataFrames: DataFrame
using CSV
Random.seed!(123);
iris = dataset("datasets", "iris")
iris = iris[shuffle(axes(iris, 1)), :]
CSV.write(joinpath(@__DIR__,"data","iris_shuffled.csv"),iris)
=#

iris     = readdlm(joinpath(@__DIR__,"data","iris_shuffled.csv"),',',skipstart=1)
x = convert(Array{Float64,2}, iris[:,1:4])
y = map(x->Dict("setosa" => 1, "versicolor" => 2, "virginica" =>3)[x],iris[:, 5])
y_oh = oneHotEncoder(y)

ntrain = Int64(round(size(x,1)*0.8))
xtrain = x[1:ntrain,:]
ytrain = y[1:ntrain]
ytrain_oh = y_oh[1:ntrain,:]
xtest = x[ntrain+1:end,:]
ytest = y[ntrain+1:end]

xtrain = scale(xtrain)
#pcaOut = pca(xtrain,error=0.01)
#(xtrain,P) = pcaOut.X, pcaOut.P
xtest = scale(xtest)
#xtest = xtest*P

l1   = DenseLayer(4,10, w=ones(10,4), wb=zeros(10),f=celu,rng=copy(TESTRNG))
l2   = DenseLayer(10,3, w=ones(3,10), wb=zeros(3),rng=copy(TESTRNG))
l3   = VectorFunctionLayer(3,f=softmax)
mynn = buildNetwork([l1,l2,l3],squaredCost,name="Multinomial logistic regression Model Sepal")
train!(mynn,xtrain,ytrain_oh,epochs=254,batchSize=8,sequential=true,verbosity=NONE,optAlg=SGD(η=t->0.001,λ=1),rng=copy(TESTRNG))
ŷtrain = predict(mynn,xtrain)
ŷtest  = predict(mynn,xtest)
trainAccuracy = accuracy(ŷtrain,ytrain,tol=1)
testAccuracy  = accuracy(ŷtest,ytest,tol=1)
@test testAccuracy >= 0.8 # set to random initialisation/training to have much better accuracy

# With ADAM
l1   = DenseLayer(4,10, w=ones(10,4), wb=zeros(10),f=celu,rng=copy(TESTRNG))
l2   = DenseLayer(10,3, w=ones(3,10), wb=zeros(3),rng=copy(TESTRNG))
l3   = VectorFunctionLayer(3,f=softmax)
mynn = buildNetwork([l1,l2,l3],squaredCost,name="Multinomial logistic regression Model Sepal")
train!(mynn,xtrain,ytrain_oh,epochs=10,batchSize=8,sequential=true,verbosity=NONE,optAlg=ADAM(η=t -> 1/(1+t), λ=0.5),rng=copy(TESTRNG))
ŷtrain = predict(mynn,xtrain)
ŷtest  = predict(mynn,xtest)
trainAccuracy = accuracy(ŷtrain,ytrain,tol=1)
testAccuracy  = accuracy(ŷtest,ytest,tol=1)
@test testAccuracy >= 1

#=
if "all" in ARGS
    # ==================================
    # NEW TEST
    # ==================================
    println("Testing colvolution layer with MINST data...")
    train_x, train_y = MNIST.traindata()
    test_x,  test_y  = MNIST.testdata()

    test = train_x[:,:,1]

    eltype(test)

    test .+ 1
end
=#

# ==================================
# NEW Test
if VERSION >= v"1.6"
    println("Testing VectorFunctionLayer with pool1d function...")
    println("Attention: this test requires at least Julia 1.6")

    x        = rand(copy(TESTRNG),300,5)
    y        = [norm(r[1:3])+2*norm(r[4:5],2) for r in eachrow(x) ]
    (N,D)    = size(x)
    l1       = DenseLayer(D,8, f=relu,rng=copy(TESTRNG))
    l2       = VectorFunctionLayer(size(l1)[2],f=(x->pool1d(x,2,f=mean)))
    l3       = DenseLayer(size(l2)[2],1,f=relu, rng=copy(TESTRNG))
    mynn     = buildNetwork([l1,l2,l3],squaredCost,name="Regression with a pooled layer")
    train!(mynn,x,y,epochs=50,verbosity=STD,rng=copy(TESTRNG))
    ŷ        = predict(mynn,x)
    mreTrain = meanRelError(ŷ,y,normRec=false)
    @test mreTrain  < 0.14
end


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

"""

    sensitivitySpecificity(ar::Array{T, 2}) where T <: Number

# Description
Calculate sensitivity and specificity from 2 x 2 array.

# Examples
```jldoctest
julia> χ = [10 40; 5 45]
julia> sensitivitySpecificity(χ)
(sensitivity = 0.6666666666666666, specificity = 0.5294117647058824)

julia> χ = [20 33; 10 37]
julia> sensitivitySpecificity(χ)
(sensitivity = 0.6666666666666666, specificity = 0.5285714285714286)
```

See also: [`predictiveValue`](@ref)
"""
function sensitivitySpecificity(ar::Array{T, 2}) where T <: Number
  if size(ar) == (2, 2)
    return (sensitivity = ar[1, 1] / ( ar[1, 1] + ar[2, 1] ), specificity = ar[2, 2] / ( ar[2, 2] + ar[1, 2] ))
  else
    @error "Matrix does not have the proper size"
  end
end

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

"""

    predictiveValue(ar::Array{T, 2}) where T <: Number

# Description
Calculate positive and negative predictive values from 2 x 2 array.

# Examples
```jldoctest
julia> χ = [10 40; 5 45]
julia> predictiveValue(χ)
(positive = 0.2, negative = 0.9)

julia> χ = [20 33; 10 37]
julia> predictiveValue(χ)
(positive = 0.37735849056603776, negative = 0.7872340425531915)
```

See also: [`sensitivitySpecificity`](@ref)
"""
function predictiveValue(ar::Array{T, 2}) where T <: Number
  if size(ar) == (2, 2)
    return (positive = ar[1, 1] / ( ar[1, 1] + ar[1, 2] ), negative = ar[2, 2] / ( ar[2, 2] + ar[2, 1] ))
  else
    @error "Matrix does not have the proper size"
  end
end

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

"""

    sensitivitySpecificity(tbVc::Array{T, 1}, labelVec::Array{T, 1}) where T <: Number

# Description
Calculate sensitivity and specificity from a `Hidden Markov model` struct.

See also: [`predictiveValue`](@ref)
"""
function sensitivitySpecificity(tbVc::Array{T, 1}, labelVec::Array{T, 1}) where T <: Number
  tbVec = copy(tbVc)

  # reassign frecuency labels
  labels = [1, 2]
  tbVec[findall(tbVec .> 1)] .= 2

  # adjust & concatenate frecuency tables
  positives = tbVec[labelVec[:, 1] .== 1] |> freqtable |> p -> convertFqDfTempl(p, templ = labels) |> p -> sort(p, rev = true)
  negatives = tbVec[labelVec[:, 1] .== 0] |> freqtable |> p -> convertFqDfTempl(p, templ = labels) |> p -> sort(p, rev = true)
  outNamedArray = [positives[:, 2] negatives[:, 2]]
  return sensitivitySpecificity(outNamedArray)
end

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

"""

    sensitivitySpecificity(tbVc::Array{T, 1}, labelMat::Matrix{T}) where T <: Number

# Description
Calculate sensitivity and specificity from a `Hidden Markov model` struct.

See also: [`predictiveValue`](@ref)
"""
function sensitivitySpecificity(tbVc::Array{T, 1}, labelMat::Matrix{T}) where T <: Number
  tbVec = copy(tbVc)

  # reassign frecuency labels
  labels = [1, 2]
  tbVec[findall(tbVec .> 1)] .= 2

  # add label columns
  labelVec = sum(labelMat, dims = 2)

  # adjust & concatenate frecuency tables
  positives = tbVec[labelVec[:, 1] .== 1] |> freqtable |> p -> convertFqDfTempl(p, templ = labels) |> p -> sort(p, rev = true)
  negatives = tbVec[labelVec[:, 1] .== 0] |> freqtable |> p -> convertFqDfTempl(p, templ = labels) |> p -> sort(p, rev = true)
  outNamedArray = [positives[:, 2] negatives[:, 2]]
  return sensitivitySpecificity(outNamedArray)
end

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

"""

    sensitivitySpecificity(ssDc::Dict{S, Tuple{Array{T, 1}, Array{Array{Float64, 1}, 1}}}, labelVec) where T <: Number where S <: String

# Description
Iterate on Dictionary and calculate sensitivity and specificity from a `Hidden Markov model` struct.

See also: [`predictiveValue`](@ref)
"""
function sensitivitySpecificity(ssDc::Dict{S, Tuple{Array{T, 1}, Array{Array{Float64, 1}, 1}}}, labelVec) where T <: Number where S <: String

  outDc = Dict{S, Array{Float64, 2}}() where S <: String
  for (k, v) in ssDc
    outSensSpec = zeros(1, 2)
    (outSensSpec[1, 1], outSensSpec[1, 2]) = sensitivitySpecificity(ssDc[k][1], labelVec)
    outDc[k] = outSensSpec
  end

  return outDc
end

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

"transform freqtable => dataframe"
function convertFqDf(fq::NamedVector{T, Array{T, 1}, Tuple{OrderedDict{T, T}}}; colnames::Array{S, 1} = ["Value", "Frecuency"]) where T <: Number where S <: String
  return DataFrames.DataFrame([names(fq)[1] fq.array], colnames)
end

"transform freqtable => dataframe template"
function convertFqDfTempl(fq::NamedVector{T, Array{T, 1}, Tuple{OrderedDict{T, T}}}; colnames::Array{S, 1} = ["Value", "Frecuency"], templ::Array{T, 1}) where T <: Number where S <: String

  fq = convertFqDf(fq)

  outDf = DataFrames.DataFrame([templ zeros(Int64, length(templ))], colnames)

  for ix in 1:size(fq, 1)
    outDf[findall(fq[ix, 1] .== outDf[:, 1]), 2] .= fq[ix, 2]
  end

  return outDf
end

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

"adjust frecuency tables for concatenation"
function stFreqTb(fTb::NamedArray{Int64, 1})
  sTb = size(fTb, 1)
  # de novo
  if sTb == 0
    fTb = [1, 2] |> freqtable |> reverse
    fTb .= 0
  # concatenate missing
  elseif sTb == 1
    added = copy(fTb)
    nPos = names(fTb)
    if sum.(nPos)[1] == 0
      NamedArrays.setnames!(added, [2], 1)
      added[1, :] .= 0
      fTb = [fTb; added]
    elseif sum.(nPos)[1] == 1
      NamedArrays.setnames!(added, [1], 1)
      added[1, :] .= 0
      fTb = [added; fTb]
    end
  # through warning
  elseif sTb > 2
    @warn "frecuency table contains more than 2 values"
  end
  return fTb
end

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


function conductive_thin_wall(;nx::Int=3, ny::Int=3, Re::Float64 = 10.0,
    Ha::Float64 = 10.0, c_w::Float64 = 0.1, α::Float64 = 1.0,
    f_u = nothing, resultsfile = nothing)

  L = 1.0
  N = Ha^2/Re
  K = Ha / (1-0.825*Ha^(-1/2)-Ha^(-1))
  ∂p∂z = -Re * K / L^3

  if f_u === nothing
    _f_u = VectorValue(0.0,0.0, -∂p∂z) * L/Re^2 # Assumed ν=L=1 by default
  else
    _f_u = f_u
  end
  g_u = VectorValue(0.0,0.0,0.0)
  g_j = VectorValue(0.0,0.0,0.0)
  B = VectorValue(0.0,1.0,0.0)

  # Discretization
  order = 2
  domain = (-1.0,1.0,-1.0,1.0,0.0,0.1)
  map(x) = VectorValue(sign(x[1])*(abs(x[1])*0.5)^0.5,
                       sign(x[2])*(abs(x[2])*0.5)^0.5, x[3]*sqrt(2)/2)*2/sqrt(2)

  dirichlet_tags_u = append!(collect(1:20),[23,24,25,26])
  dirichlet_tags_j = append!(collect(1:20),[25,26])
  neumann_j = [23,24]

  partition=(nx,ny,3)
  model = CartesianDiscreteModel(domain,partition;
    isperiodic=(false,false,true), map=map)

  labels = get_face_labeling(model)
  add_tag_from_tags!(labels,"dirichlet_u",dirichlet_tags_u)
  add_tag_from_tags!(labels,"dirichlet_j",dirichlet_tags_j)
  add_tag_from_tags!(labels,"neumann_j",neumann_j)

  Vu = FESpace(model, ReferenceFE(lagrangian,VectorValue{3,Float64},order);
      conformity=:H1, dirichlet_tags="dirichlet_u")

  Vp = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:P);
      conformity=:L2, constraint=:zeromean)

  Vj = FESpace(model, ReferenceFE(raviart_thomas,Float64,order-1);
      conformity=:Hdiv, dirichlet_tags="dirichlet_j")

  Vφ = FESpace(model, ReferenceFE(lagrangian,Float64,order-1,space=:Q);
      conformity=:L2)

  U = TrialFESpace(Vu,g_u)
  P = TrialFESpace(Vp)
  J = TrialFESpace(Vj,g_j)
  Φ = TrialFESpace(Vφ)

  Y = MultiFieldFESpace([Vu, Vp, Vj, Vφ])
  X = MultiFieldFESpace([U, P, J, Φ])

  trian = Triangulation(model)
  degree = 2*(order)
  dΩ = Measure(trian,degree)

  btrian_j = BoundaryTriangulation(model,tags=["neumann_j"])
  dΓ = Measure(btrian_j,degree)
  n = get_normal_vector(btrian_j)

  res_Ω(x,y) = InductionlessMHD.dimensionless_residual(x, y, Re, N, B, f_u, dΩ)
  jac_Ω(x,dx,y) = InductionlessMHD.dimensionless_jacobian(x, dx, y, Re, N, B, dΩ)
  res_Γ(x,y) = InductionlessMHD.dimensionless_conducting_wall(x, y, n, c_w, dΓ, α=α)
  jac_Γ(x,dx,y) = InductionlessMHD.dimensionless_conducting_wall(dx, y, n, c_w, dΓ, α=α)

  res(x,y) = res_Ω(x,y) + res_Γ(x,y)
  jac(x,dx,y) = jac_Ω(x,dx,y) + jac_Γ(x,dx,y)
  op  = FEOperator(res,jac,X,Y)

  nls = NLSolver(;
    show_trace=true, method=:newton, linesearch=BackTracking())
  solver = FESolver(nls)

  xh = solve(solver,op)

  if resultsfile != nothing
    uh, ph, jh, φh = xh
    writevtk(trian, resultsfile,
      cellfields=["uh"=>uh, "ph"=>ph, "jh"=>jh, "φh"=>φh])
  end

  (xh, trian, dΩ)
end


using Cat

const Tensor = Array{Float64}
@category TensorNetwork

@arrow TensorNetwork TSum :: (Tensor, Tensor) --> Tensor
@arrow TensorNetwork TProd :: (Tensor, Tensor) --> Tensor
@arrow TensorNetwork Contract :: (Tensor, Int, Int) --> Tensor

function example()
    g = tensor(2) # metric
    R = tensor(4) # Riemann curvature tensor
    T = tensor(2) # Stress-energy tensor

    Ric = Contract(R, 1, 3) # Ricci tensor
    Rsc = Contract(Ric, 1, 2) # Ricci scalar
    EFE = Ric - 0.5*Rsc*g - (8*π)*T
end

@alias TensorNetwork TSum Base.:+ 2
@alias TensorNetwork TProd Base.:* 2
Base.:*(a::Float64, b::TensorNetwork.Arrow) = TProd(a*ones(), b)
Base.:*(a::TensorNetwork.Arrow, b::Float64) = TProd(b*ones(), a)

tensor(nd::Int64) = TensorNetwork.Identity{Tensor{nd}}()

@interpretation Einsum (=>) TensorNetwork begin
    letters::Dict{TensorNetwork.Arrow, Char}
end


# Here we solve the equation on a unit square, with a uniform internal source.
# The strong form of the (linear) heat equation is given by
#
# ```math
#  -\nabla \cdot (k \nabla u) = f  \quad x \in \Omega,
# ```
#
# where $u$ is the unknown temperature field, $k$ the heat conductivity,
# $f$ the heat source and $\Omega$ the domain. For simplicity we set $f = 1$
# and $k = 1$. We will consider homogeneous Dirichlet boundary conditions such that
# ```math
# u(x) = 0 \quad x \in \partial \Omega,
# ```
# where $\partial \Omega$ denotes the boundary of $\Omega$.
#
# The HDG weak form is given by
# ```math
#(\vec{\sigma}_{h},\vec{v_{h}})_{\tau_{h}}-(u_{h},\nabla_{h}\cdot\vec{v}_{h})_{\tau_{h}}+\left\langle \hat{u}_{h},\vec{v}_{h}\cdot\vec{n}_{h}\right\rangle _{\partial\tau_{h}} & =0\quad\forall\vec{v}_{h}\in\vec{V}_{h},\\
#(w_{h},\nabla_{h}\cdot\vec{\sigma}_{h})_{\tau_{h}}+\left\langle (\hat{\sigma}_{h}-\vec{\sigma}_{h})\cdot\vec{n},w_{h}\right\rangle _{\partial\tau_{h}} & =(f,w_{h})_{\tau_{h}}\quad\forall w_{h}\in W_{h},\\
#\left\langle \vec{\sigma}_{h},\xi_{h}\right\rangle _{\partial\tau_{h}\backslash\Gamma} & =0\quad\forall\xi_{h}\in M_{h}^{o},\\
#\left\langle \hat{u}_{h},\xi_{h}\right\rangle _{\Gamma} & =\left\langle g,\xi_{h}\right\rangle _{\Gamma}\quad\forall\xi_{h}\in M_{h}^{\partial},\\
#\hat{\sigma}_{h}\cdot\vec{n} & =\vec{\sigma}_{h}\cdot\vec{n}+\tau(u_{h}-\hat{u}_{h}).
# ```
#-

using HDiscontinuousGalerkin
using Tensors
using BlockArrays
using LinearAlgebra

# Load mesh
#To load mesh from triangle
#root_file = "mesh/figure.1"
#@time mesh = parse_mesh_triangle(root_file)

#or use internal mesh
mesh = rectangle_mesh(TriangleCell, (10,10), Vec{2}((0.0,0.0)), Vec{2}((1.0,1.0)))

# ### Trial and test functions
dim = 2
finiteElement = GenericFiniteElement(Dubiner{dim,RefTetrahedron,1}())
Wh = ScalarFunctionSpace(mesh, finiteElement)
Vh = VectorFunctionSpace(mesh, finiteElement)
Mh = ScalarTraceFunctionSpace(Wh, GenericFiniteElement(Legendre{dim-1,RefTetrahedron,1}()))

# Variables
û_h = TrialFunction(Mh)
σ_h = TrialFunction(Vh)
u_h = TrialFunction(Wh)

# ### Boundary conditions
dbc = Dirichlet(û_h, mesh, "boundary", x -> 0)

# RHS function
f(x::Vec{dim}) = 2*π^2*sin(π*x[1])*sin(π*x[2])
# ### Assembling the linear system
# Now we have all the pieces needed to assemble the linear system, $K û = f$.
function doassemble(Vh, Wh, Mh, τ = 1.0)
    # Allocate Matrices
    n_basefuncs = getnbasefunctions(Vh)
    n_basefuncs_s = getnbasefunctions(Wh)
    n_basefuncs_t = getnbasefunctions(Mh)
    Ae = zeros(n_basefuncs, n_basefuncs)
    Be = zeros(n_basefuncs, n_basefuncs_s)
    Ce = zeros(n_basefuncs_s, n_basefuncs_s)
    Ee = zeros(n_basefuncs,3*n_basefuncs_t)
    Fe = zeros(n_basefuncs_s, 3*n_basefuncs_t)
    be = zeros(n_basefuncs_s,1)
    He = zeros(3*n_basefuncs_t, 3*n_basefuncs_t)

    # create a matrix assembler and rhs vector
    assembler = start_assemble(getnfaces(mesh)*n_basefuncs_t)
    rhs = zeros(Float64, getnfaces(mesh)*n_basefuncs_t)
    K_element = Array{AbstractMatrix{Float64}}(undef, getncells(mesh))
    b_element = Array{AbstractVector{Float64}}(undef, getncells(mesh))

    @inbounds for (cell_idx, cell) in enumerate(CellIterator(mesh))
        reinit!(Wh, cell)
        reinit!(Vh, cell)
        fill!(Ae, 0)
        fill!(Be, 0)
        fill!(Ce, 0)
        fill!(Ee, 0)
        fill!(Fe, 0)
        fill!(He, 0)
        fill!(be, 0)
        #Cell integrals
        for q_point in 1:getnquadpoints(Vh)
            dΩ = getdetJdV(Vh, q_point)
            for i in 1:n_basefuncs
                vh  = shape_value(Vh, q_point, i)
                div_vh = shape_divergence(Vh, q_point, i)
                for j in 1:n_basefuncs
                    σ = shape_value(Vh, q_point, j)
                    # Integral σ⋅v dΩ
                    Ae[i, j] += (σ ⋅ vh) * dΩ
                end
                for j in 1:n_basefuncs_s
                    u = shape_value(Wh, q_point, j)
                    # Integral u*∇⋅v dΩ
                    Be[i,j] += (u*div_vh) * dΩ
                end
            end
        end
        #RHS
        for q_point in 1:getnquadpoints(Wh)
            dΩ = getdetJdV(Wh, q_point)
            for i in 1:n_basefuncs_s
                w  = shape_value(Wh, q_point, i)
                fh = function_value(f, Wh, cell, q_point)
                # Integral f*u dΩ
                be[i] += fh*w*dΩ
            end
        end
        #Face integrals
        for face_idx in 1:getnfaces(cell)
            for q_point in 1:getnfacequadpoints(Wh)
                dS = getfacedetJdS(Wh, face_idx, q_point)
                orientation = face_orientation(cell, face_idx)
                for i in 1:n_basefuncs_s
                    w = shape_value(Wh, face_idx, q_point, i)
                    for j in 1:n_basefuncs_s
                        u = shape_value(Wh, face_idx, q_point, j)
                        # Integral_∂T τ*u ⋅ w dS
                        Ce[i,j] += τ*(u*w) * dS
                    end
                    w = shape_value(Wh, face_idx, q_point, i, orientation)
                    for j in 1:n_basefuncs_t
                        û = shape_value(Mh, q_point, j)
                        # Integral_∂T τ*û*w  dS
                        Fe[i,n_basefuncs_t*(face_idx-1)+j] += (τ*(û*w)) * dS
                    end
                end
                dS = getfacedetJdS(Mh, face_idx, q_point)
                n = get_normal(Vh, face_idx)
                for i in 1:n_basefuncs
                    v = shape_value(Vh, face_idx, q_point, i, orientation)
                    for j in 1:n_basefuncs_t
                        û = shape_value(Mh, q_point, j)
                        # Integral_∂T û(v⋅n)  dS
                        Ee[i,n_basefuncs_t*(face_idx-1)+j] += (û*(v⋅n)) * dS
                    end
                end
                for i in 1:n_basefuncs_t
                    μ = shape_value(Mh, q_point, i)
                    for j in 1:n_basefuncs_t
                        û = shape_value(Mh, q_point, j)
                        # Integral_∂T û*μ  dS
                        He[n_basefuncs_t*(face_idx-1)+i,n_basefuncs_t*(face_idx-1)+j] += (û*μ) * dS
                    end
                end
            end
        end
        #Assamble Ke
        Me = BlockArray{Float64}(undef, [n_basefuncs,n_basefuncs_s], [n_basefuncs,n_basefuncs_s])
        setblock!(Me, Ae, 1, 1)
        setblock!(Me, Be', 2, 1)
        setblock!(Me, -Be, 1, 2)
        setblock!(Me, Ce, 2, 2)
        Mei = factorize(Array(Me))
        EFe = BlockArray{Float64}(undef, [n_basefuncs,n_basefuncs_s], [3*n_basefuncs_t])
        setblock!(EFe, -Ee, 1, 1)
        setblock!(EFe, Fe, 2, 1)
        Ge = BlockArray{Float64}(undef, [n_basefuncs,n_basefuncs_s], [3*n_basefuncs_t])
        setblock!(Ge, Ee, 1, 1)
        setblock!(Ge, Fe, 2, 1)
        Ge = Ge'
        Bte = BlockArray{Float64}(undef, [n_basefuncs,n_basefuncs_s],[1])
        setblock!(Bte, zeros(n_basefuncs,1),1,1)
        setblock!(Bte, be,2,1)
        K_element[cell_idx]=Mei\EFe
        Ate = Ge*K_element[cell_idx]-He
        b_element[cell_idx]=(Mei\Bte)[:,1]
        bte = -Ge*b_element[cell_idx]
        #Assemble
        gdof = Vector{Int}()
        for fidx in mesh.cells[cell_idx].faces
            for j in 1:n_basefuncs_t
                push!(gdof, fidx*n_basefuncs_t-(n_basefuncs_t-j))
            end
        end
        assemble!(assembler, gdof, Ate)
        assemble!(rhs, gdof, bte[:,1])
    end
    return end_assemble(assembler), rhs, K_element, b_element
end

### Solution of the system
K, b, K_e, b_e = doassemble(Vh,Wh,Mh);

# To account for the boundary conditions we use the `apply!` function.
# This modifies elements in `K` and `f` respectively, such that
# we can get the correct solution vector `u` by using `\`.
apply!(K,b,dbc)
û = K \ b;
#Now we recover original variables from skeleton û
function get_uσ!(σ_h,u_h,û_h,û, K_e, b_e, mesh)
    n_cells = getncells(mesh)
    n_basefuncs = getnbasefunctions(σ_h)
    n_basefuncs_t = getnbasefunctions(Mh)
    for cell_idx in 1:getncells(mesh)
        #get dofs
        û_e = Vector{Float64}()
        for (k,face) in enumerate(mesh.cells[cell_idx].faces)
            push!(û_e, û[n_basefuncs_t*face-1:n_basefuncs_t*face]...)
            û_h.m_values[cell_idx,:,k] = û[n_basefuncs_t*face-1:n_basefuncs_t*face]
        end
        dof = K_e[cell_idx]*û_e+b_e[cell_idx]
        σ_h.m_values[cell_idx,:] = dof[1:n_basefuncs]
        u_h.m_values[cell_idx,:] = dof[n_basefuncs+1:end]
    end
end

get_uσ!(σ_h, u_h,û_h,û, K_e, b_e, mesh)
#Compute errors
u_ex(x::Vec{dim}) = sin(π*x[1])*sin(π*x[2])
Etu_h = errornorm(u_h, u_ex)
Etu_h <= 0.00006

### Plot Solution
# TODO: Turn into a recipe
using Makie, CairoMakie
#AbstractPlotting.__init__()
using ColorSchemes
connectivity = getcells_matrix(mesh)
coordinates = get_vertices_matrix(mesh)

# We avg since u_h is discontinuous
color = nodal_avg(u_h)
scene = Scene(resolution = (400, 200), colorrange = (0.0, 1.0), colormap = ColorSchemes.viridis.colors)
#poly!(scene, coordinates, connectivity, color = color, strokecolor = (:black, 0.6), strokewidth = 1)
for row in 1:size(connectivity)[1]
	#read coordinates
	points = node(:poly, Point2f0[coordinates[i,:] for i in connectivity[row,:]])
	colors = get(ColorSchemes.viridis, [color[i] for i in connectivity[row,:]])
    #colors = get(ColorSchemes.viridis, [value(u_h, node, row) for node in connectivity[row,:]])
	poly!(scene, points, strokewidth = 1, strokecolor = (:black,0.6), color = colors, show_axis = true, scale_plot = false)
end
display(scene)


"""
The Dist object contains information for convolving random variables.
* `lb`: lower bound of density
* `ub`: upper bound of density
* `n_points`: number of discrete points used in transform
* `x`: real values associated with densities
* `density`: densities resulting from Convolution
* `interp_pdf`: an object used to interpolate exact densities from x and density
```@julia
model = Normal(0,1) + Normal(0,1)
convolve!(model)

julia> pdf(model, 0.0)
0.28216441829987216

julia> pdf(Normal(0,sqrt(2)), 0.0)
0.28209479177387814
```
"""
mutable struct Convolution{T}
    model::T
    lb::Float64
    ub::Float64
    n_points::Int64
    x::Vector{Float64}
    density::Vector{Float64}
    interp_pdf::Spline1D
end

function Convolution(;model, lb=-3.0, ub=3.0, n_points = 2^9, x=fill(0.0, 0),
        density=fill(0.0, 0), interp_pdf=Spline1D(zeros(4), zeros(4)))
    Convolution(model, lb, ub, n_points, x, density, interp_pdf)
end

Broadcast.broadcastable(c::Convolution) = Ref(c)

function +(d::Distribution...)
    convolution = Convolution(model=d)
    return convolution
end

function analytic(D::Convolution)
    N = D.n_points;  model = D.model
    lb = D.lb;  ub = D.ub
    i = range(0, length=N); dx = (ub-lb)/N
    x = lb .+ i*dx; dt = 2*π/(N*dx)
    c = -N/2*dt; t = c .+ i*dt; f = fill(1im+1.0, N)
    for m in model
        @. f = f*cf(m, t)
    end
    X = @. exp(-(0+1im)*i*dt*lb)*f
    fft!(X)
    density = @. dt/(2*π)*exp(-(0+1im)*c*x)*X
    # Use abs()to haldle very small negative values, e.g. -2.50e-12
    D.density = abs.(real(density))
    D.x = x
    D.interp_pdf = Spline1D(D.x, D.density)
    return nothing
end

function numeric(D::Convolution)
   n_points = D.n_points;  model = D.model
   lb = D.lb;  ub = D.ub
   Ndist = length(model)
   stepsize = (ub-lb)/n_points
   x = range(lb, stop=ub, length=n_points)
   d = pdf.(model[1], x)
   mn = Int(round(1.5*n_points)) + 1
   st = Int(round(2*n_points))
   mx = n_points ÷ 2
   v = [fill(0.0, mx); d ;fill(0.0, mx)]
   Nv = length(v)
   d1 = pdf.(model[2], x)
   v1 = [fill(0.0, mx); d1; fill(0.0, mx)]
   m = fft(v).*fft(v1)
   cv = real(ifft(m))*Nv*stepsize/(2*n_points)
   cv = abs.(cv[[mn:st;1:mx]])
   cnt = 2
   while cnt < Ndist
       cnt += 1
       v = [fill(0.0, mx); cv; fill(0.0, mx)]
       d1 = pdf.(model[cnt], x)
       v1 = [fill(0.0, mx); d1 ;fill(0.0, mx)]
       m = fft(v).*fft(v1)
       cv = real(ifft(m))*Nv*stepsize/(2*n_points)
       cv = abs.(cv[[mn:st;1:mx]])
    end
    D.density = cv
    D.x = x
    D.interp_pdf = Spline1D(D.x, D.density)
    return nothing
end

"""
Convolves a set of random variables.
```@julia
model = Normal(0,1) + Normal(0,1)
convolve!(model)

julia> pdf(model, 0.0)
0.28216441829987216

julia> pdf(Normal(0,sqrt(2)), 0.0)
0.28209479177387814
```
"""
function convolve!(D::Convolution)
    if all(x->applicable(cf, (x,1.0)...), D.model)
        # Convolve with characteristic function if possible
        analytic(D)
    else
        numeric(D)
    end
    return nothing
end

 function rand(model::Convolution, N::Int64)
     rts = fill(0.0, N)
     dist = model.model
     for d in dist
         rts .+= rand(d, N)
     end
     return rts
 end

Gradient(d::Convolution, x::Float64) = d.interpGrad[x]
pdf(d::Convolution, x::Float64) = abs(d.interp_pdf(x))
logpdf(d::Convolution, x::Float64) = log(pdf(d, x))

"""
Normal approximation to the sum of Uniform rvs
* 'scaling' = (2/3): default scaling used for standard deviation σ
* 'args': a list of NamedTuples for component distributions, e.g. cr =(μ=.05,N=2),...

```@julia

julia> μ,σ = convolve_normal(cr=(μ=.05,N=2),vis=(μ=.085,N=1))
(0.185, 0.021278575558006173)
```
"""
function convolve_normal(;scaling=2/3, args...)
    # uniform factor
    fact1 = 1/12
    μ,σ = 0.0,0.0
    for (k,v) in pairs(args)
        μ += prod(v)
        σ += fact1.*v.N*(scaling*v.μ)^2
    end
    σ = sqrt(σ)
    return μ,σ
end


using JuMP, EAGO, Gurobi

m = Model(with_optimizer(EAGO.Optimizer, relaxed_optimizer = Gurobi.Optimizer(OutputFlag=0)))

# ----- Variables ----- #
x_Idx = Any[2, 3, 4, 5, 6, 7]
@variable(m, x[x_Idx])
JuMP.set_lower_bound(x[2], 1.0e-7)
JuMP.set_upper_bound(x[2], 0.2)
JuMP.set_lower_bound(x[3], 1.0e-7)
JuMP.set_upper_bound(x[3], 0.2)
JuMP.set_lower_bound(x[4], 1.0e-7)
JuMP.set_upper_bound(x[4], 0.4)
JuMP.set_lower_bound(x[5], 1.0e-7)
JuMP.set_upper_bound(x[5], 0.4)
JuMP.set_lower_bound(x[6], 1.0e-7)
JuMP.set_upper_bound(x[6], 0.4)
JuMP.set_lower_bound(x[7], 1.0e-7)
JuMP.set_upper_bound(x[7], 0.4)

# ----- Constraints ----- #
@constraint(m, e2, x[2]+x[3] == 0.2)
@constraint(m, e3, x[4]+x[5] == 0.4)
@constraint(m, e4, x[6]+x[7] == 0.4)

# ----- Objective ----- #
@NLobjective(m, Min, ((15.3261663216011*x[2]+23.2043471859416*x[4]+6.69678129464404*x[6])*log(2.1055*x[2]+3.1878*x[4]+0.92*x[6])-2.46348749603266*x[2]-4.33155441248417*x[4]-0.626542690017204*x[6]+6.4661663216011*x[2]*log(x[2]/(2.1055*x[2]+3.1878*x[4]+0.92*x[6]))+12.2043471859416*x[4]*log(x[4]/(2.1055*x[2]+3.1878*x[4]+0.92*x[6]))+0.696781294644034*x[6]*log(x[6]/(2.1055*x[2]+3.1878*x[4]+0.92*x[6]))+9.86*x[2]*log(x[2]/(1.972*x[2]+2.4*x[4]+1.4*x[6]))+12*x[4]*log(x[4]/(1.972*x[2]+2.4*x[4]+1.4*x[6]))+7*x[6]*log(x[6]/(1.972*x[2]+2.4*x[4]+1.4*x[6]))+(1.972*x[2]+2.4*x[4]+1.4*x[6])*log(1.972*x[2]+2.4*x[4]+1.4*x[6])+1.972*x[2]*log(x[2]/(1.972*x[2]+0.283910843616504*x[4]+3.02002220174195*x[6]))+2.4*x[4]*log(x[4]/(1.45991339466884*x[2]+2.4*x[4]+0.415073537580851*x[6]))+1.4*x[6]*log(x[6]/(0.602183324335333*x[2]+0.115623371371275*x[4]+1.4*x[6]))+(15.3261663216011*x[3]+23.2043471859416*x[5]+6.69678129464404*x[7])*log(2.1055*x[3]+3.1878*x[5]+0.92*x[7])-2.46348749603266*x[3]-4.33155441248417*x[5]-0.626542690017204*x[7]+6.4661663216011*x[3]*log(x[3]/(2.1055*x[3]+3.1878*x[5]+0.92*x[7]))+12.2043471859416*x[5]*log(x[5]/(2.1055*x[3]+3.1878*x[5]+0.92*x[7]))+0.696781294644034*x[7]*log(x[7]/(2.1055*x[3]+3.1878*x[5]+0.92*x[7]))+9.86*x[3]*log(x[3]/(1.972*x[3]+2.4*x[5]+1.4*x[7]))+12*x[5]*log(x[5]/(1.972*x[3]+2.4*x[5]+1.4*x[7]))+7*x[7]*log(x[7]/(1.972*x[3]+2.4*x[5]+1.4*x[7]))+(1.972*x[3]+2.4*x[5]+1.4*x[7])*log(1.972*x[3]+2.4*x[5]+1.4*x[7])+1.972*x[3]*log(x[3]/(1.972*x[3]+0.283910843616504*x[5]+3.02002220174195*x[7]))+2.4*x[5]*log(x[5]/(1.45991339466884*x[3]+2.4*x[5]+0.415073537580851*x[7]))+1.4*x[7]*log(x[7]/(0.602183324335333*x[3]+0.115623371371275*x[5]+1.4*x[7]))-17.2981663216011*x[2]*log(x[2])-25.6043471859416*x[4]*log(x[4])-8.09678129464404*x[6]*log(x[6])-17.2981663216011*x[3]*log(x[3])-25.6043471859416*x[5]*log(x[5])-8.09678129464404*x[7]*log(x[7])))

JuMP.optimize!(m)
run_time = backend(m).optimizer.model.optimizer._run_time
println("run time: $run_time")


using PolarimetryCalibration, Plots, LaTeXStrings

σ = 0.01
PAsrc = 36.0
PArange = 0:0.2:180
const (I,Q,U,V)=(1,2,3,4)

stokes = [(müller_rec()*
           müller_sky(PAaz*π/180)*
           lin_pol_source(pa = PAsrc*π/180))[i] + σ*randn()
          for PAaz = PArange, i in 1:4]
plot(PArange, stokes[:,V], label=L"ideal", xlabel = "PA", ylabel = "Stokes Q")

stokes = [(müller_rec(ΔG = 0.1)*
           müller_sky(PAaz*π/180)*
           lin_pol_source(pa = PAsrc*π/180))[i] + σ*randn()
          for PAaz = PArange, i in 1:4]
plot!(PArange, stokes[:,V], label=L"\Delta G = 0.1")

stokes = [(müller_rec(ψ = 52*π/180)*
           müller_sky(PAaz*π/180)*
           lin_pol_source(pa = PAsrc*π/180))[i] + σ*randn()
          for PAaz = PArange, i in 1:4]
plot!(PArange, stokes[:,V], label=L"$\psi = 52$ $^\ocirc$")


using LinearAlgebra

Δx = 0.04
Lx = 4.0
xlim = (-Lx/2,Lx/2)
ylim = (-Lx/2,Lx/2)
g = PhysicalGrid(xlim,ylim,Δx)

RadC = 1.0
Δs = 1.4*cellsize(g)
body = Circle(RadC,Δs)
X = VectorData(collect(body))

w = Nodes(Dual,size(g))
ϕ = Nodes(Primal,size(g))
q = Edges(Primal,size(g))
f = ScalarData(X)

angs(n) = range(0,2π,length=n+1)[1:n]

@testset "Forming a cache" begin
  scache1 = SurfaceScalarCache(body,g,scaling=GridScaling)
  scache2 = SurfaceScalarCache(X,g,scaling=GridScaling)
  @test normals(scache1) == normals(scache2)
  @test areas(scache1) == areas(scache2)
  @test points(scache1) == points(scache2)
end

@testset "Scalar surface ops" begin
  scache = SurfaceScalarCache(body,g,scaling=GridScaling)

  gd = similar_grid(scache)
  @test typeof(gd) == typeof(scache.gdata_cache)

  sd = similar_surface(scache)
  @test typeof(sd) == typeof(scache.sdata_cache)

  f .= sin.(angs(length(body)) .- π/4)
  regularize_normal!(q,f,scache)
  normal_interpolate!(f,q,scache)
  @test maximum(f) ≈ 11 atol = 1e-0
  @test minimum(f) ≈ -11 atol = 1e-0

  A = create_CLinvCT(scache,scale=cellsize(g))

  @test maximum(abs.(eigvals(A))) ≈ 0.2 atol = 1e-1

  A = create_GLinvD(scache,scale=cellsize(g))
  @test maximum(abs.(eigvals(A))) ≈ 0.45 atol = 1e-1

  A = create_nRTRn(scache)
  @test maximum(svdvals(A)) ≈ 11 atol = 1e-0

  # Create additional regularization and interpolation ops
  w = ones_gridcurl(scache)
  Rn = RegularizationMatrix(scache,scache.sdata_cache,w)
  En = InterpolationMatrix(scache,w,scache.sdata_cache)

  f .= En*w

end

dq = EdgeGradient(Primal,size(g))
vs = VectorData(X)

@testset "Vector surface ops" begin
  vcache = SurfaceVectorCache(body,g,scaling=GridScaling)

  vs.u .= sin.(angs(length(body)) .- π/4)

  regularize_normal!(dq,vs,vcache)
  normal_interpolate!(vs,dq,vcache)
  @test minimum(vs.u) ≈ -39 atol = 1e-0
  @test maximum(vs.u) ≈ 39 atol = 1e-0

  surface_divergence!(q,vs,vcache)
  surface_grad!(vs,q,vcache)

  A = create_GLinvD(vcache,scale=cellsize(g))
  @test maximum(abs.(eigvals(A))) ≈ 1.7 atol = 1e-1

end

@testset "Masks" begin

  scache = SurfaceScalarCache(body,g,scaling=GridScaling)

  inner = mask(scache)

  oc = similar(ϕ)
  oc .= 1
  @test abs(dot(oc,inner,g) - π*RadC^2) < 2e-3

  outer = complementary_mask(scache)

end

@testset "Grid operations" begin
  scache = SurfaceScalarCache(body,g,scaling=GridScaling)

  q .= randn(size(q))
  cv_cache = ConvectiveDerivativeCache(EdgeGradient(Primal,q))
  qdq = similar(q)
  qdq .= 0.0
  convective_derivative!(qdq,q,scache,cv_cache)

  qdq2 = convective_derivative(q,scache)

  @test qdq == qdq2
end

@testset "Problem specification" begin

  prob = BasicScalarILMProblem(g,body,scaling=GridScaling)
  sys = ImmersedLayers.__init(prob)

  scache = SurfaceScalarCache(body,g,scaling=GridScaling)

  @test typeof(sys.base_cache) == typeof(scache)

  vprob = BasicVectorILMProblem(g,body,scaling=GridScaling)
  vsys = ImmersedLayers.__init(vprob)

  vcache = SurfaceVectorCache(body,g,scaling=GridScaling)

  @test typeof(vsys.base_cache) == typeof(vcache)

  p = rand()
  prob = BasicScalarILMProblem(g,body,scaling=GridScaling,phys_params=p)
  sys = ImmersedLayers.__init(prob)
  @test sys.base_cache.phys_params == p

  Re = rand()
  p = Dict("Re"=>Re)
  prob = BasicScalarILMProblem(g,body,scaling=GridScaling,phys_params=p)
  sys = ImmersedLayers.__init(prob)
  @test sys.base_cache.phys_params["Re"] == Re

  my_bc_func(x) = x
  prob = BasicScalarILMProblem(g,body,scaling=GridScaling,bc=my_bc_func)
  sys = ImmersedLayers.__init(prob)
  @test sys.base_cache.bc == my_bc_func

  my_f_func(x) = x
  prob = BasicScalarILMProblem(g,body,scaling=GridScaling,f_funcs=my_f_func)
  sys = ImmersedLayers.__init(prob)
  @test sys.base_cache.f_funcs == my_f_func

end


# Borrowed from http://www.mathworks.com/help/signal/ref/upfirdn.html

using  Winston
import Multirate

sampleRate    = 48000
interpolation = 147
decimation    = 160
ratio         = interpolation  // decimation
Nt            = 24*interpolation

h = Multirate.firdes( Nt, 0.5/interpolation, Multirate.kaiser, beta = 7.8562  )

n = 0:10239
x = sin(2*pi*1e3/sampleRate*n)
y = Multirate.filt( h, x, ratio )

tOriginal = n[1:49]./sampleRate
xOriginal = x[1:49]
tResamp   = n[1:45]./(sampleRate*ratio)
xResamp   = y[13:57]

display( stem( tOriginal, xOriginal, "b." ) )
hold( true )
display( stem( tResamp, xResamp, "r." ) )
hold( false )


function convectionCenteredX(u::AbstractGridGhost, vloc, i, j)
    0.5 * u[i, j] * (u[i+1, j] - u[i-1, j]) / u.grid.dx
        + 0.5 * vloc * (u[i, j+1] - u[i, j-1]) / u.grid.dy
end

function convectionCenteredY(v::AbstractGridGhost, uloc, i, j)
    0.5 * uloc * (v[i+1, j] - v[i-1, j]) / v.grid.dx
        + 0.5 * v[i, j] * (v[i, j+1] - v[i, j-1]) / v.grid.dy
end

function laplacianCentered(g::AbstractGrid, i, j)
    dx2 = 1 / (g.grid.dx * g.grid.dx)
    dy2 = 1 / (g.grid.dy * g.grid.dy)
    c = g[i, j]
    dx2 * (g[i+1, j] - 2. * c + g[i-1, j]) + dy2 * (g[i, j+1] - 2. * c + g[i, j-1])
end

function pressureX(p::AbstractGrid, i, j)
    (p[i, j] - p[i-1, j]) / p.grid.dx
end

function pressureY(p::AbstractGrid, i, j)
    (p[i, j] - p[i, j-1]) / p.grid.dy
end

function fillghost!(g::UGrid, f1, f2)
    ghost1 = g.ghost[1]
    ghost2 = g.ghost[2]
    for i in eachindex(ghost1)
        ghost1[i] = f1(g[i, 1], g[i, 2], g[i, 3])
        ghost2[i] = f2(g[i, end], g[i, end-1], g[i, end-2])
    end
end

function fillghost!(g::VGrid, f1, f2)
    ghost1 = g.ghost[1]
    ghost2 = g.ghost[2]
    for i in eachindex(ghost1)
        ghost1[i] = f1(g[1, i], g[2, i], g[1, i])
        ghost2[i] = f2(g[end, i], g[end-1, i], g[end-2, i])
    end
end

function noslip(arg::Vararg{Real})
    u1, u2, u3 = arg
    -0.2 * (u3 - 5. * u2 + 15. * u1)
end

function naturalbound(arg::Vararg{Real})
    u1 = arg[1]
    u1
end

function veldiv(u::UGrid, v::VGrid, i, j)
    (u[i+1, j] - u[i, j]) / u.grid.dx + (v[i, j+1] - v[i, j]) / v.grid.dy
end

function step_AB2!(sim::Simulation, u_n_1::UGrid, v_n_1::VGrid)
    u_n = copy(sim.u)
    v_n = copy(sim.v)
    nu = getnu(sim.params, sim.u.grid)
    xu = getxs(sim.u)
    yu = getys(sim.u)

    xv = getxs(sim.v)
    yv = getys(sim.v)
    for i in 2:size(sim.u)[1]-1
        for j in 1:size(sim.u)[2]
    # for ci in eachindex(inner_u)
            # (i, j) = Tuple(ci)
            vloc = projectionU(v_n, i, j)
            vloc_n_1 = projectionU(v_n_1, i, j)
            # f1 = - 0.5 * (3. * convectionCenteredX(u_n, vloc, i, j) - convectionCenteredX(u_n_1, vloc_n_1, i, j))
            f1 = - AB2(convectionCenteredX(u_n, vloc, i, j), convectionCenteredX(u_n_1, vloc_n_1, i, j))
            f2 = - pressureX(sim.p, i, j)
            f3 = nu*laplacianCentered(u_n, i, j)

            with_logger(sim.logger) do
                @debug "f1: $f1 - f2: $f2 - f3: $f3"
            end

            mask = inbump(sim.canal, xu[i], yu[j])
            # sim.u[ci] = mask ? (f1 + f2 + f3) / (1 + sim.params.dtdτ) : (f1 + f2 + f3)
            fn1 = u_n[i, j] + sim.dt * (f1 + f2 + f3)
            if mask
                sim.u[i, j] = fn1 / (1 + sim.params.dtdτ * ramp_khi(sim))
            else
                sim.u[i, j] = fn1
            end
            # sim.u[i, j] = mask ? (f1 + f2 + f3) / (1 + sim.params.dtdτ) : (f1 + f2 + f3)
        end
    end

    # for ci in eachindex(inner_v)
    for i in 1:size(sim.v)[1]
        for j in 2:size(sim.v)[2]-1
        # (i, j) = Tuple(ci)
            uloc = projectionV(u_n, i, j)
            uloc_n_1 = projectionV(u_n_1, i, j)
            f1 = - AB2(convectionCenteredY(v_n, uloc, i, j), convectionCenteredY(v_n_1, uloc_n_1, i, j))
            f2 = - pressureY(sim.p, i, j)
            f3 = nu*laplacianCentered(v_n, i, j)

            mask = inbump(sim.canal, xv[i], yv[j])

            fn1 = v_n[i, j] + sim.dt * (f1 + f2 + f3)
            if mask
                sim.v[i, j] = fn1 / (1 + sim.params.dtdτ * ramp_khi(sim))
            else
                sim.v[i, j] = fn1
            end
            # sim.v[ci] = mask ? (f1 + f2 + f3) / (1 + sim.params.dtdτ) : (f1 + f2 + f3)
            # sim.v[i, j] = mask ? (f1 + f2 + f3) / (1 + sim.params.dtdτ) : (f1 + f2 + f3)
        end
    end

    boundaries!(sim)
end

function step_euler!(sim::Simulation)
    u_n = copy(sim.u)
    v_n = copy(sim.v)
    nu = getnu(sim.params, sim.u.grid)
    xu = getxs(sim.u)
    yu = getys(sim.u)

    xv = getxs(sim.v)
    yv = getys(sim.v)
    for i in 2:size(sim.u)[1]-1
        for j in 1:size(sim.u)[2]
    # for ci in eachindex(inner_u)
            # (i, j) = Tuple(ci)
            vloc = projectionU(v_n, i, j)
            f1 = - convectionCenteredX(u_n, vloc, i, j)
            f2 = - pressureX(sim.p, i, j)
            f3 = nu*laplacianCentered(u_n, i, j)

            with_logger(sim.logger) do
                @debug "f1: $f1 - f2: $f2 - f3: $f3"
            end

            mask = inbump(sim.canal, xu[i], yu[j])
            # sim.u[ci] = mask ? (f1 + f2 + f3) / (1 + sim.params.dtdτ) : (f1 + f2 + f3)
            fn1 = u_n[i, j] + sim.dt * (f1 + f2 + f3)
            if mask
                sim.u[i, j] = fn1 / (1 + sim.params.dtdτ * ramp_khi(sim))
            else
                sim.u[i, j] = fn1
            end
            # sim.u[i, j] = mask ? (f1 + f2 + f3) / (1 + sim.params.dtdτ) : (f1 + f2 + f3)
        end
    end

    # for ci in eachindex(inner_v)
    for i in 1:size(sim.v)[1]
        for j in 2:size(sim.v)[2]-1
        # (i, j) = Tuple(ci)
            uloc = projectionV(u_n, i, j)
            f1 = - convectionCenteredY(v_n, uloc, i, j)
            f2 = - pressureY(sim.p, i, j)
            f3 = nu*laplacianCentered(v_n, i, j)

            mask = inbump(sim.canal, xv[i], yv[j])

            fn1 = v_n[i, j] + sim.dt * (f1 + f2 + f3)
            if mask
                sim.v[i, j] = fn1 / (1 + sim.params.dtdτ * ramp_khi(sim))
            else
                sim.v[i, j] = fn1
            end
            # sim.v[ci] = mask ? (f1 + f2 + f3) / (1 + sim.params.dtdτ) : (f1 + f2 + f3)
            # sim.v[i, j] = mask ? (f1 + f2 + f3) / (1 + sim.params.dtdτ) : (f1 + f2 + f3)
        end
    end

    boundaries!(sim)
end

function rightboundary!(sim::Simulation)
    yu = getys(sim.u)
    for col in eachcol(sim.u)
        j = col.indices[2]
        col[end] = col[end] - sim.dt * poiseuille(sim, yu[j]) * (col[end] - col[end-1]) / sim.u.grid.dx
    end
end

function boundaries!(sim::Simulation)
    fillghost!(sim.u, noslip, noslip)
    fillghost!(sim.v, noslip, naturalbound)
    rightboundary!(sim)
end

# TO BE VERIFIED, SEEMS NOT GOOD
function pgauss(p::PGrid, i, j, dx2, dy2)
    right = i != p.grid.m ? p[i+1, j] : -p[i, j]
    left = i != 1 ? p[i-1, j] : p[i, j]
    top = j != p.grid.n ? p[i, j+1] : p[i, j] # DP is 0 at walls
    bot = j != 1 ? p[i, j-1] : p[i, j] # DP is 0 at walls

    (right + left) / dx2 + (top + bot) / dy2
end

function gauss_seidel!(p::PGrid, u::UGrid, v::VGrid, α, dt)
    dx2 = u.grid.dx * u.grid.dx
    dy2 = u.grid.dy * u.grid.dy
    coef = 0.5 * dx2 * dy2 / (dx2 + dy2)
    # innerp = @view p[2:end-1, 2:end-1]
    for ci in eachindex(p)
        i, j = Tuple(ci)
        div = - veldiv(u, v, i, j) / dt
        phit = pgauss(p, i, j, dx2, dy2)
        phistar = coef * ( div + phit )
        p[ci] = α * phistar + (1 - α) * p[ci]
    end
end
gauss_seidel!(sim::Simulation) = gauss_seidel!(sim.p, sim.u, sim.v, sim.params.alpha, sim.dt)

function residual(p::PGrid, u::UGrid, v::VGrid, dt)
    R = 0
    # SHOULD TAKE INTO ACCOUNT BOUNDARIES
    is, js = inrange(p)
    for i in is
        for j in js
            # i, j = Tuple(ci)
            lap = laplacianCentered(p, i, j)
            div = veldiv(u, v, i, j)
            R += (lap - div / dt)^2
        end
    end
    R * p.grid.dx * p.grid.dy
end
function residual(sim::Simulation)
    Rs = residual(sim.p, sim.u, sim.v, sim.dt)
    grid = sim.u.grid
    sim.dt * grid.D / getum(sim.params, grid) * sqrt( 1 / (grid.L * grid.D) * Rs )
end

function step_poisson!(sim::Simulation, tol = 1e-3)
    e = 1000.
    n = 0
    maxn = 1000
    while (abs(e) > tol) && (n < maxn)
        gauss_seidel!(sim)
        e = residual(sim)
        n += 1
    end

    if n == maxn
        with_logger(sim.logger) do
            @warn "Gauss Seidel exceeded $maxn iterations! Aborted"
        end
    end

    with_logger(sim.logger) do
        @info "Poisson equation solved after $n iterations"
    end
end

function poisson_project!(u::UGrid, v::VGrid, p::PGrid, dt)
    ustar = copy(u)
    vstar = copy(v)

    is, js = inrange(u)
    for i in is
        for j in js
            u[i, j] = ustar[i, j] - dt / u.grid.dx * (p[i-1, j] - p[i, j])
        end
    end

    is, js = inrange(v)
    for i in is
        for j in js
            v[i, j] = vstar[i, j] - dt / u.grid.dy * (p[i, j] - p[i, j-1])
        end
    end
end
poisson_project!(sim::Simulation) = poisson_project!(sim.u, sim.v, sim.p, sim.dt)
# function euler!(g::AbstractGridGhost, Δt, f::Function)
#     for ei in eachindex(g)
#         i,j = Tuple(ei)
#         g[ei] = g[ei] + Δt * f(g, i, j)
#     end
# end

function euler(u_n, Δt, f_n)
    u_n + Δt * f_n
end

# function AB2(u_n, Δt, f_n, f_n_1)
#     u_n + 0.5 * Δt * (-f_n_1 + 3. * f_n)
# end

function AB2(f_n, f_n_1)
    0.5 * (-f_n_1 + 3. * f_n)
end

function advectionDiffusionX(g::AbstractGridGhost, i, j)

end

function ramp_khi(ramp::Int, n::Int)
    coef = n / ramp
    coef > 1. ? 1. : coef
end

ramp_khi(sim::Simulation) = ramp_khi(sim.params.ramp, sim.step)


###here is Pody layer
const tanh_a = (0, 1, 0, 1/3, 0, 2/15)
Aaf(n) = tanh_a[1 : n]

pody(x) = max(-1, min(1, evalpoly(x, (0, 1, 0, 1/3))))

struct PD{F, M<:AbstractMatrix, B}
    W::M
    b::B
    σ::F
    function PD(W, b)
        return new{typeof(pody), typeof(W), typeof(b)}(W, b, pody)
    end
end

function (layer::PD)(x::AbstractVecOrMat)
    W, b, σ = layer.W, layer.b, layer.σ
    σ.(W * x .+ b)
end

function Base.show(io::IO, a::PD)
    o, i = size(a.W)
    print("PD(IO:$i => $o, pody:[0, 1, 0, 1/3] * x)")
end

Flux.trainable(a::PD) = a.a, a.W
Flux.@functor PD
###end of Pody layer

###here is Data Keeper
struct DataKeeper
    data::Array{Tuple, 1}
    function DataKeeper()
        new(Array{Tuple}(undef, 0))
    end
end

function Base.push!(DK::DataKeeper, x::Tuple)
    if length(DK.data) == 100
        popfirst!(DK.data)
    end
    push!(DK.data, x)
end
###end of Data Keeper


function test_diophantine_solve_diophantine_quadratic()
    soltype = Array{AbstractDiophantineSolutions{DiophantineSolutionXY{Int}},1}

    # linear
    sol = solve(diophantine_equation_quadratic_xy(cx=3, c0=-9))
    @test Set(sol) == Set([diophantine_onex_anyy(3)])
    @test length(sol) == 1
    @test typeof(sol) == soltype

    # simple hyperbolic
    sol = solve(diophantine_equation_quadratic_xy(cxy=4, cx=2, cy=8, c0=4))
    @test Set(sol) == Set([diophantine_onex_anyy(-2)])
    @test length(sol) == 1
    @test typeof(sol) == soltype

    sol = solve(diophantine_equation_quadratic_xy(cxy=4, cx=12, cy=8, c0=24))
    @test Set(sol) == Set([diophantine_onex_anyy(-2), diophantine_anyx_oney(-3)])
    @test length(sol) == 2
    @test typeof(sol) == soltype

    sol = solve(diophantine_equation_quadratic_xy(cxy=2, cx=5, cy=56, c0=7))
    @test Set(sol) == Set([diophantine_solutions((-27, 64), (-29, -69), (-21, 7), (-35, -12), (-9, 1), (-47, -6), (105, -2), (-161, -3))])
    @test length(sol) == 1
    @test typeof(sol) == soltype

    # elliptical
    sol = solve(diophantine_equation_quadratic_xy(cx²=42, cxy=8, cy²=15, cx=23, cy=17, c0=-4915))
    @test Set(sol) == Set([diophantine_solutions((-11, -1))])
    @test length(sol) == 1
    @test typeof(sol) == soltype

    # parabolic
    sol = solve(diophantine_equation_quadratic_xy(cx²=1, cxy=2, cy²=1, cx=2, cy=2, c0=1))
    @test Set(sol) == Set([diophantine_linearx_lineary(1, 0, -1, -1)])
    @test length(sol) == 1
    @test typeof(sol) == soltype

    sol = solve(diophantine_equation_quadratic_xy(cx²=1, cxy=2, cy²=1, cx=6, cy=6, c0=5))
    @test Set(sol) == Set([diophantine_linearx_lineary(1, 0, -1, -1), diophantine_linearx_lineary(1, 0, -1, -5)])
    @test length(sol) == 2
    @test typeof(sol) == soltype

    sol = solve(diophantine_equation_quadratic_xy(cx²=8, cxy=-24, cy²=18, cx=5, cy=7, c0=16))
    @test Set(sol) == Set([diophantine_quadraticx_quadraticy(-174, 17, -2, -116, 21, -2), diophantine_quadraticx_quadraticy(-174, 41, -4, -116, 37, -4)])
    @test length(sol) == 2
    @test typeof(sol) == soltype

    # hyperbolic (WIP)
    sol = solve(diophantine_equation_quadratic_xy(cx²=1, cxy=3, cy²=2))
    @test Set(sol) == Set([diophantine_solutions((0, 0)), diophantine_linearx_lineary(2, 0, -1, 0), diophantine_linearx_lineary(1, 0, -1, 0)])
    @test length(sol) == 3
    @test typeof(sol) == soltype

    sol = solve(diophantine_equation_quadratic_xy(cx²=1, cxy=3, cy²=2, c0=10))
    @test Set(sol) == Set([diophantine_solutions((-21, 11), (21, -11), (-12, 7), (12, -7), (-9, 7), (9, -7), (-12, 11), (12, -11))])
    @test length(sol) == 1
    @test typeof(sol) == soltype
end


struct WendlandC6{T} <: SPHKernel
    n_neighbours::Int64
    norm_1D::T
    norm_2D::T
    norm_3D::T
end

"""
    WendlandC6(T::DataType=Float64, n_neighbours::Integer=295)

Set up a `WendlandC6` kernel for a given DataType `T`.
"""
WendlandC6(T::DataType=Float64, n_neighbours::Integer=295) = WendlandC6{T}(n_neighbours, T(55.0/32.0), T(78.0/(7.0*π)), T(1365.0/(64.0*π)))

"""
    kernel_value_1D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

Evaluate WendlandC6 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_value_1D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

    if u < 1
        n = kernel.norm_1D * h_inv
        t1 = 1 - u
        t6 = t1*t1*t1*t1*t1*t1
        u2 = u*u
        return ( t6 * t1 * (1 + 7u + 19u2 + 21u2 * u)) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    kernel_deriv_1D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

Evaluate the derivative of the WendlandC6 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_deriv_1D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

     if u < 1
        n = kernel.norm_1D * h_inv^2
        t1 = 1 - u
        t6 = t1*t1*t1*t1*t1*t1
        u2 = u*u
        return ( -6t6 * u * (35u2 + 18u + 3 )) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    bias_correction_1D(kernel::WendlandC6, density::Real, m::Real, h_inv::Real)

Corrects the density estimate for the kernel bias. See Dehnen&Aly 2012, eq. 18+19.
"""
function bias_correction_1D(kernel::WendlandC6{T}, density::Real, m::Real, h_inv::Real) where T

    n = kernel.norm_1D * h_inv^3
    wc_correction = T(0.0116) * ( kernel.n_neighbours * T(0.01) )^T(-2.236) * m * n

    if wc_correction < T(0.2)*density
        density -= wc_correction
    end

    return density |> T
end

"""
    (kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

Evaluate WendlandC6 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_value_2D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

    if u < 1
        n = kernel.norm_2D * h_inv^2
        u_m1 = 1 - u
        u_m1 = u_m1 * u_m1  # (1.0 - u)^2
        u_m1 = u_m1 * u_m1  # (1.0 - u)^4
        u_m1 = u_m1 * u_m1  # (1.0 - u)^8
        u2 = u*u
        return ( u_m1 * ( 1 + 8u + 25u2 + 32u2*u )) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    kernel_deriv_2D(kernel::WendlandC6, u::Real, h_inv::Real)

Evaluate the derivative of the WendlandC6 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_deriv_2D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

    if u < 1
        n = kernel.norm_2D * h_inv^3
        u_m1 = 1 - u
        u_m1_7 = u_m1 * u_m1 * u_m1 * u_m1 * u_m1 * u_m1 * u_m1
        return ( -22u_m1_7 * u * ( 16*u^2 + 7u + 1 )) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    bias_correction_2D(kernel::WendlandC6{T}, density::Real, m::Real, h_inv::Real) where T

Corrects the density estimate for the kernel bias. See Dehnen&Aly 2012, eq. 18+19.
"""
function bias_correction_2D(kernel::WendlandC6{T}, density::Real, m::Real, h_inv::Real) where T

    n = kernel.norm_2D * h_inv^3
    wc_correction = T(0.0116) * ( kernel.n_neighbours * T(0.01) )^T(-2.236) * m * n

    if wc_correction < T(0.2)*density
        density -= wc_correction
    end

    return density |> T
end

"""
    kernel_value_3D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

Evaluate WendlandC6 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_value_3D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

    if u < 1
        n = kernel.norm_3D * h_inv^3
        u_m1 = 1 - u
        u_m1 = u_m1 * u_m1  # (1.0 - u)^2
        u_m1 = u_m1 * u_m1  # (1.0 - u)^4
        u_m1 = u_m1 * u_m1  # (1.0 - u)^8
        u2 = u*u
        return ( u_m1 * ( 1 + 8u + 25u2 + 32u2*u )) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    kernel_deriv_3D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

Evaluate the derivative of the WendlandC6 spline at position ``u = \\frac{x}{h}``.
"""
function kernel_deriv_3D(kernel::WendlandC6{T}, u::Real, h_inv::Real) where T

    if u < 1
        n = kernel.norm_3D * h_inv^2 * h_inv^2
        u_m1 = 1 - u
        u_m1_7 = u_m1 * u_m1 * u_m1 * u_m1 * u_m1 * u_m1 * u_m1
        return ( -22u_m1_7 * u * ( 16u^2 + 7u + 1 )) * n |> T
    else
        return 0.0 |> T
    end

end

"""
    bias_correction_3D(kernel::WendlandC6, density::Real, m::Real, h_inv::Real)

Corrects the density estimate for the kernel bias. See Dehnen&Aly 2012, eq. 18+19.
"""
function bias_correction_3D(kernel::WendlandC6{T}, density::Real, m::Real, h_inv::Real) where T

    n = kernel.norm_3D * h_inv^3
    wc_correction::T = T(0.0116) * ( kernel.n_neighbours * T(0.01) )^T(-2.236) * m * n

    if wc_correction < T(0.2)*density
        density -= wc_correction
    end

    return density |> T
end