# # Approximations of the bilinear term

# Ultimately we want to evolve, in time, an equation of the form # $$
#   \omega_t = G(\omega).
# $$
# For the 2D NSE, $G$ is of the form
# $$
#  G(\omega) = g + \nu \Delta \omega + B(\omega),
# $$
# where
# $$
#   B(\omega) = \left(\partial_x^2 - \partial_y^2\right)(uv) + \partial_{xy}\left(v^2 - u^2\right),
# $$
# with $\mathbf{u} = (u(x, y), v(x, y))$ being the velocity field associated with the vorticity, i.e.  $\omega = \boldsymbol{\nabla} \times \mathbf{u} = v_x - u_y$.

# In spectral space, $\Delta \omega$ is just a diagonal operator, while $g$ is a given scalar field, which are both easy to compute. The most computationally demanding term is $B(\omega)$.

# Here, we test different ways of approximating $B(\omega)$.

# The packages we are gonna need.

using FFTW
using Plots
using LinearAlgebra: mul!
using Test
using BenchmarkTools
using Random

@info "Threads: $(FFTW.nthreads())"

# ## Operators

# We first define a method to build the various operators acting in spectral space:

function get_operators(N, κ₀)
    ## Differentiation in spectral space
    Dx_hat = im * κ₀ * [
        ifelse(k1 ≤ div(N, 2) + 1, k1 - 1, k1 - 1 - N) for k2 in 1:div(N, 2)+1, k1 in 1:N
    ]
    Dy_hat = im * κ₀ * [k2 - 1 for k2 in 1:div(N, 2)+1, k1 in 1:N]

    Delta_hat = - κ₀^2 * [
        ifelse(k1 ≤ div(N, 2) + 1, (k1 - 1)^2 + (k2 - 1)^2, (k1 - 1 - N)^2 + (k2 - 1)^2)
        for k2 in 1:div(N, 2)+1, k1 in 1:N
    ]

    ## For the Basdevant formulation
    DxsqDysq_hat = Dx_hat.^2 .- Dy_hat.^2
    Dxy_hat = Dx_hat .* Dy_hat

    ## Recovering of the velocity field from the vorticity
    Hu_hat = - Dy_hat ./ Delta_hat
    Hu_hat[1, 1] = 0.0
    Hv_hat = Dx_hat ./ Delta_hat
    Hv_hat[1, 1] = 0.0

    return Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat, DxsqDysq_hat, Dxy_hat
end

# ## Generation of a scalar field

# Method to generate a scalar field from a list of wavenumbers and amplitudes:

function field_from_modes(L, N, modes::Matrix{<:Integer}, amps::Matrix{<:Real})
    κ₀ = 2π/L
    x = y = (L/N):(L/N):L
    field = sum(
        [
            2κ₀^2 * (k[1]^2 + k[2]^2) * (
                a[1] * cos.(κ₀ * (k[1] * one.(y) * x' + k[2] * y * one.(x)'))
                - a[2] * sin.(κ₀ * (k[1] * one.(y) * x' + k[2] * y * one.(x)'))
            )
            for (k, a) in zip(eachrow(modes), eachrow(amps))
        ]
    )
    return field
end

# Method to randomly generate a scalar field from a given random number generator and a given number wavenumbers to be excited:

function field_from_modes(rng::AbstractRNG, L, N, num_modes::Int)

    modes = rand(rng, 1:div(N,10), num_modes, 2)
    amps = rand(rng, num_modes, 2)

    field = field_from_modes(L, N, modes, amps)

    return field
end

# Method to randomly generate a scalar field from a given number wavenumbers to be excited:

field_from_modes(L, N, num_modes::Int) = field_from_modes(Xoshiro(), L, N, num_modes)

# ## Methods for computing the bilinear term

function bilinear_naive(vort_hat, params)
    operators, vars = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat = operators
    N, Nsub = vars

    u_hat = Hu_hat .* vort_hat
    v_hat = Hv_hat .* vort_hat
    u = irfft(u_hat, N)
    v = irfft(v_hat, N)
    vort = irfft(vort_hat, N)
    wu_hat = rfft(vort .* u)
    wv_hat = rfft(vort .* v)

    bilin_hat = Dx_hat .* wu_hat .+ Dy_hat .* wv_hat
    ## dealiasing
    bilin_hat[div(Nsub,2) + 1:end, :] .= 0.0im
    bilin_hat[:, div(Nsub,2) + 1:div(N,2) + div(Nsub,2)] .= 0.0im
    return bilin_hat
end

#

function bilinear_auxs!(bilin_hat, vort_hat, params)
    operators, vars, auxs = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat = operators
    N, Nsub = vars
    u_hat, v_hat, u, v, w, wu, wv, wu_hat, wv_hat = auxs

    u_hat .= Hu_hat .* vort_hat
    v_hat .= Hv_hat .* vort_hat
    u .= irfft(u_hat, N)
    v .= irfft(v_hat, N)
    w .= irfft(vort_hat, N)
    wu .= w .* u
    wv .= w .* v
    wu_hat .= rfft(wu)
    wv_hat .= rfft(wv)

    bilin_hat .= Dx_hat .* wu_hat .+ Dy_hat .* wv_hat
    ## dealiasing
    bilin_hat[div(Nsub,2) + 1:end, :] .= 0.0im
    bilin_hat[:, div(Nsub,2) + 1:div(N,2) + div(Nsub,2)] .= 0.0im

    return bilin_hat
end

#

function bilinear_plan_brdcst!(bilin_hat, vort_hat, params)
    operators, vars, auxs, plans = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat = operators
    N, Nsub = vars
    u_hat, v_hat, u, v, w, wu, wv, wu_hat, wv_hat = auxs
    plan, plan_inv = plans

    u_hat .= Hu_hat .* vort_hat
    v_hat .= Hv_hat .* vort_hat
    mul!(u, plan_inv, u_hat)
    mul!(v, plan_inv, v_hat)
    ## mul! on inverse plan mutates last argument, but we need to preserve vort_hat
    ## so, reuse v_hat for that:
    v_hat .= vort_hat
    mul!(w, plan_inv, v_hat)
    wu .= w .* u
    wv .= w .* v

    mul!(wu_hat, plan, wu)
    mul!(wv_hat, plan, wv)

    bilin_hat .= Dx_hat .* wu_hat .+ Dy_hat .* wv_hat
    ## dealiasing
    bilin_hat[div(Nsub,2) + 1:end, :] .= 0.0im
    bilin_hat[:, div(Nsub,2) + 1:div(N,2) + div(Nsub,2)] .= 0.0im

    return bilin_hat
end

#

function bilinear_Basdevant_brdcst(vort_hat, params)
    operators, vars = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat = operators
    N, Nsub = vars

    u_hat = Hu_hat .* vort_hat
    v_hat = Hv_hat .* vort_hat
    u = irfft(u_hat, N)
    v = irfft(v_hat, N)
    uv_hat = rfft(u .* v)
    v2u2_hat = rfft(v.^2 - u.^2)

    bilin_hat = DxsqDysq_hat .* uv_hat .+ Dxy_hat .* v2u2_hat
    ## dealiasing
    bilin_hat[div(Nsub,2) + 1:end, :] .= 0.0im
    bilin_hat[:, div(Nsub,2) + 1:div(N,2) + div(Nsub,2)] .= 0.0im

    return bilin_hat
end

#

function bilinear_Basdevant_plan_brdcst!(bilin_hat, vort_hat, params)
    operators, vars, auxs, plans = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat = operators
    N, Nsub = vars
    u_hat, v_hat, u, v, w, uv, v2u2, uv_hat, v2u2_hat = auxs
    plan, plan_inv = plans

    u_hat .= Hu_hat .* vort_hat
    v_hat .= Hv_hat .* vort_hat

    mul!(u, plan_inv, u_hat)
    mul!(v, plan_inv, v_hat)

    uv .= u .* v
    v2u2 .= v.^2 .- u.^2
    mul!(uv_hat, plan, uv)
    mul!(v2u2_hat, plan, v2u2)

    bilin_hat .= DxsqDysq_hat .* uv_hat .+ Dxy_hat .* v2u2_hat

    ## dealiasing
    bilin_hat[div(Nsub,2) + 1:end, :] .= 0.0im
    bilin_hat[:, div(Nsub,2) + 1:div(N,2) + div(Nsub,2)] .= 0.0im

    return bilin_hat
end

#

function bilinear_Basdevant_plan_loop!(bilin_hat, vort_hat, params)
    operators, vars, auxs, plans = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat = operators
    N, Nsub = vars
    u_hat, v_hat, u, v, w, uv, v2u2, uv_hat, v2u2_hat = auxs
    plan, plan_inv = plans

    @inbounds for k in eachindex(u_hat)
        u_hat[k] = Hu_hat[k] * vort_hat[k]
        v_hat[k] = Hv_hat[k] * vort_hat[k]
    end

    mul!(u, plan_inv, u_hat)
    mul!(v, plan_inv, v_hat)

    @inbounds for k in eachindex(u)
        uv[k] = u[k] * v[k]
        v2u2[k] = v[k]^2 - u[k]^2
    end

    mul!(uv_hat, plan, uv)
    mul!(v2u2_hat, plan, v2u2)

    @inbounds for k in eachindex(vort_hat)
        bilin_hat[k] = DxsqDysq_hat[k] * uv_hat[k] + Dxy_hat[k] * v2u2_hat[k]
    end
    ## dealiasing
    @inbounds for j in 1:N, i in div(Nsub,2) + 1:div(N, 2) + 1
        bilin_hat[i, j] = 0.0im
    end
    @inbounds for j in div(Nsub,2) + 1:div(N,2) + div(Nsub,2), i in 1:div(N, 2) + 1
        bilin_hat[i, j] = 0.0im
    end

    return bilin_hat
end

#

function bilinear_Basdevant_plan_doubleloop!(bilin_hat, vort_hat, params)
    operators, vars, auxs, plans = params

    Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat,
        DxsqDysq_hat, Dxy_hat = operators
    N, Nsub = vars
    u_hat, v_hat, u, v, w, uv, v2u2, uv_hat, v2u2_hat = auxs
    plan, plan_inv = plans

    M, N = size(vort_hat)

    @inbounds for j in 1:N, i in 1:M
        u_hat[i, j] = Hu_hat[i, j] * vort_hat[i, j]
        v_hat[i, j] = Hv_hat[i, j] * vort_hat[i, j]
    end

    mul!(u, plan_inv, u_hat)
    mul!(v, plan_inv, v_hat)

    @inbounds for j in 1:N, i in 1:N
        uv[i, j] = u[i, j] * v[i, j]
        v2u2[i, j] = v[i, j]^2 - u[i, j]^2
    end

    mul!(uv_hat, plan, uv)
    mul!(v2u2_hat, plan, v2u2)

    @inbounds for j in 1:N, i in 1:M
        bilin_hat[i, j] = DxsqDysq_hat[i, j] * uv_hat[i, j] + Dxy_hat[i, j] * v2u2_hat[i, j]
    end
    ## dealiasing
    @inbounds for j in 1:N, i in div(Nsub,2) + 1:M
        bilin_hat[i, j] = 0.0im
    end
    @inbounds for j in div(Nsub,2) + 1:div(N,2) + div(Nsub,2), i in 1:M
        bilin_hat[i, j] = 0.0im
    end

    return bilin_hat
end

# ## The spatial domain and its discretization

L = 2π
κ₀ = 2π/L
N = 256
Nsub = 162
x = y = (L/N):(L/N):L

# ## Preparation for tests

Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat, DxsqDysq_hat, Dxy_hat = get_operators(N, κ₀)

vort = field_from_modes(L, N, 48)
vort_hat = rfft(vort)

u_hat = similar(vort_hat)
v_hat = similar(vort_hat)
u = similar(vort)
v = similar(vort)
w = similar(vort)
uv = similar(vort)
v2u2 = similar(vort)
uv_hat = similar(vort_hat)
v2u2_hat = similar(vort_hat)
bilin_hat = similar(vort_hat)

plan = plan_rfft(vort, flags=FFTW.MEASURE)
plan_inv = plan_irfft(vort_hat, N, flags=FFTW.MEASURE)

operators = Dx_hat, Dy_hat, Delta_hat, Hu_hat, Hv_hat, DxsqDysq_hat, Dxy_hat
vars = N, Nsub
auxs = u_hat, v_hat, u, v, w, uv, v2u2, uv_hat, v2u2_hat
plans = plan, plan_inv

# ## Check results

# As a first test, we check that all implementations return about the same vector

@testset "same result bilinear terms" begin

    bilin_hat_naive = bilinear_naive(
        vort_hat,
        (operators, vars)
    )

    bilinear_hat_brdcst = similar(vort_hat)
    bilinear_auxs!(bilinear_hat_brdcst, vort_hat, (operators, vars, auxs))

    bilin_hat_plan_brdcst = similar(vort_hat)
    bilinear_plan_brdcst!(
        bilin_hat_plan_brdcst,
        vort_hat,
        (operators, vars, auxs, plans)
    )

    bilin_hat_Basdevant_brdcst = bilinear_Basdevant_brdcst(
        vort_hat,
        (operators, vars)
    )

    bilin_hat_Basdevant_plan_brdcst = similar(vort_hat)
    bilinear_Basdevant_plan_brdcst!(
        bilin_hat_Basdevant_plan_brdcst,
        vort_hat,
        (operators, vars, auxs, plans)
    )

    bilin_hat_Basdevant_plan_loop = similar(vort_hat)
    bilinear_Basdevant_plan_loop!(
        bilin_hat_Basdevant_plan_loop,
        vort_hat,
        (operators, vars, auxs, plans)
    )

    bilin_hat_Basdevant_plan_doubleloop = similar(vort_hat)
    bilinear_Basdevant_plan_doubleloop!(
        bilin_hat_Basdevant_plan_doubleloop,
        vort_hat,
        (operators, vars, auxs, plans)
    )

    @test bilin_hat_naive ≈ bilinear_hat_brdcst
    @test bilinear_hat_brdcst ≈ bilin_hat_plan_brdcst
    @test bilin_hat_plan_brdcst ≈ bilin_hat_Basdevant_brdcst
    @test bilin_hat_Basdevant_brdcst ≈ bilin_hat_Basdevant_plan_brdcst
    @test bilin_hat_Basdevant_plan_brdcst ≈ bilin_hat_Basdevant_plan_loop
    @test bilin_hat_Basdevant_plan_loop ≈ bilin_hat_Basdevant_plan_doubleloop
end
nothing

# ## Orthogonality

# As a second test, we check the orthogonality condition of the bilinear term.

# Due to the periodicity and the divergence-free conditions, we should have $(B(\omega), \omega)_{L^2} = 0$. We compute $B(\omega)$ in spectral space and transform it to physical space in order to check that.

bilin_hat_naive = bilinear_naive(vort_hat, (operators, vars))
bilin_naive = irfft(bilin_hat_naive, N)
nothing

# We compute the inner product by summing over the nodal values and multiplying that by the area element $\Delta x \Delta y = (L/N)^2$. Notice computing the inner product as such has an inherent numerical error, which is also of the order of the area element, so we don't expect the output to be of machine precision.

sum(bilin_naive .* vort) * (L / N)^2

# In spectral space, due to Parseval's identity, this translates to the following (except we need to correct for some indices...)

M = iseven(N) ? div(N, 2) : div(N, 2) + 1

(
    2 * sum(bilin_hat_naive[2:M,:] .* conj.(vort_hat[2:M,:]))
    + sum(bilin_hat_naive[1,:] .* conj.(vort_hat[1,:]))
    + sum(bilin_hat_naive[M+1:end,:] .* conj.(vort_hat[M+1:end,:]))
) / N^2 * (L / N)^2

# Hmm, something is wrong...

# ## Benchmark

@info "bilinear_naive"
@btime bilinear_naive(vh, p) setup = (vh = copy($vort_hat); p = $(operators, vars));

@info "bilinear_auxs!"
@btime bilinear_auxs!(bh, vh, p) setup = (bh = similar($vort_hat); vh = copy($vort_hat); p = $(operators, vars, auxs, plans));

@info "bilinear_Basdevant_brdcst"
@btime bilinear_Basdevant_brdcst(vh, p) setup = (vh = copy($vort_hat); p = $(operators, vars));

@info "bilinear_plan_brdcst!"
@btime bilinear_plan_brdcst!(bh, vh, p) setup = (bh = similar($vort_hat); vh = copy($vort_hat); p = $(operators, vars, auxs, plans));

@info "bilin_hat_Basdevant_plan_brdcst!"
@btime bilinear_Basdevant_plan_brdcst!(bh, vh, p) setup = (bh = similar($vort_hat); vh = copy($vort_hat); p = $(operators, vars, auxs, plans));

@info "bilinear_Basdevant_plan_loop!"
@btime bilinear_Basdevant_plan_loop!(bh, vh, p) setup = (bh = similar($vort_hat); vh = copy($vort_hat); p = $(operators, vars, auxs, plans));

@info "bilinear_Basdevant_plan_doubleloop!"
@btime bilinear_Basdevant_plan_doubleloop!(bh, vh, p) setup = (bh = similar($vort_hat); vh = copy($vort_hat); p = $(operators, vars, auxs, plans));

# ## Conclusions

# Clearly, `bilinear_naive` is pretty bad, in great part because of the use of the allocating transforms `rfft` and `irfft`.

# The method `bilinear_auxs!` saves some allocations but is pretty slow as well, due to also using `rfft` and `irfft`.

# The method `bilinear_Basdevant_brdcst` saves some time by using the Basdevant formulation, but is still quite slow because of `rfft` and `irfft`.

# All the others are much faster, due to using the plans for the transforms and auxiliary functions to avoid allocation of temporary variable.

# Of the nonallocating ones, the ones that use the Basdevant formulation have a clear advantage. Indeed, computing the direct and inverse Fourier transform are the most costly operation here, so reducing it from five to four makes a significant difference.

# Of these nonallocating ones that use the Basdevant formulation, the two that use explicit loops perform better than the broadcast one, with a slight advantage for the one that uses a single loop with `eachindex` instead of a double loop.

# If one does [@code_warntype](https://docs.julialang.org/en/v1/manual/performance-tips/index.html#man-code-warntype-1) on them, one sees that most are clean, but there are some yellow warnings with type `Union{Nothing, Tuple{Int64, Int64}}` precisely on the two fastest ones. Yellow warnings, fortunately, tend to be harmless, as is the case here. In this case, the warning is due to the way the [iterator interface](https://docs.julialang.org/en/v1/manual/interfaces/#Interfaces) is implemented, which returns `nothing` when the end of the iteration is reached. One can get rid of it with `while` loops, but it won't improve the performance.


# # Spatio-Temporal Prediction Examples
# In this page we are simply running files from the
# `examples` folder of the `TimeseriesPrediction` package.

# This is how you can (programmatically) find this folder:
# ```julia
# using TimeseriesPrediction
# exdir = dirname(dirname(pathof(TimeseriesPrediction)))*"/examples"
# ```

# ## Temporal Prediction: Kuramoto-Sivashinsky
# *(this requires `FFTW` to be installed)*

include("1Dfield_temporalprediction.jl")

# ## Cross Prediction: Barkley Model

include("2Dfield_crossprediction.jl")

# ## Temporal Prediction: Periodic Nonlinear Barkley Model

include("2Dfield_temporalprediction.jl")


function part1(xs)
    sum(diff(xs) .> 0)
end

function valid_conv(xs, w)
    n = length(xs)
    m = length(w)
    rw = reverse(w)

    [sum(xs[i:i+m-1] .* rw) for i=1:(n-m+1)]
end

function part2(xs)
    w = repeat(1:1, 3)
    part1(valid_conv(xs, w))
end

xs = open(ARGS[1]) do fp
    map(x -> parse(Int, x), split(read(fp, String)))
end

println("Part 1: ", part1(xs))
println("Part 2: ", part2(xs))


#################### Raftery and Lewis Diagnostic ####################

"""
    rafterydiag(x::AbstractVector{<:Real}; q, r, s, eps, range)

Compute the Raftery and Lewis diagnostic.
"""
function rafterydiag(
    x::AbstractVector{<:Real}; q=0.025, r=0.005, s=0.95, eps=0.001, range=1:length(x)
)
    nx = length(x)
    phi = sqrt(2.0) * SpecialFunctions.erfinv(s)
    nmin = ceil(Int, q * (1.0 - q) * (phi / r)^2)
    if nmin > nx
        @warn "At least $nmin samples are needed for specified q, r, and s"
        kthin = -1
        burnin = total = NaN
    else
        dichot = Int[(x .<= StatsBase.quantile(x, q))...]
        kthin = 0
        bic = 1.0
        local test , ntest
        while bic >= 0.0
            kthin += 1
            test = dichot[1:kthin:nx]
            ntest = length(test)
            temp = test[1:(ntest - 2)] + 2 * test[2:(ntest - 1)] + 4 * test[3:ntest]
            trantest = reshape(StatsBase.counts(temp, 0:7), 2, 2, 2)
            g2 = 0.0
            for i1 in 1:2, i2 in 1:2, i3 in 1:2
                tt = trantest[i1, i2, i3]
                if tt > 0
                    fitted =
                        sum(trantest[:, i2, i3]) * sum(trantest[i1, i2, :]) /
                        sum(trantest[:, i2, :])
                    g2 += 2.0 * tt * log(tt / fitted)
                end
            end
            bic = g2 - 2.0 * log(ntest - 2.0)
        end

        tranfinal = StatsBase.counts(test[1:(ntest - 1)] + 2 * test[2:ntest], 0:3)
        alpha = tranfinal[3] / (tranfinal[1] + tranfinal[3])
        beta = tranfinal[2] / (tranfinal[2] + tranfinal[4])
        kthin *= step(range)
        m = log(eps * (alpha + beta) / max(alpha, beta)) / log(abs(1.0 - alpha - beta))
        burnin = kthin * ceil(m) + first(range) - 1
        n = ((2.0 - alpha - beta) * alpha * beta * phi^2) / (r^2 * (alpha + beta)^3)
        keep = kthin * ceil(n)
        total = burnin + keep
    end
    return (
        thinning=kthin, burnin=burnin, total=total, nmin=nmin, dependencefactor=total / nmin
    )
end


abstract type GeneralizedKroneckerProduct{T} <: AbstractMatrix{T} end

Base.eltype(K::GeneralizedKroneckerProduct{T}) where {T} = T

abstract type AbstractKroneckerProduct{T} <: GeneralizedKroneckerProduct{T} end

Base.IndexStyle(::Type{<:GeneralizedKroneckerProduct}) = IndexCartesian()

"""
    collect!(C::AbstractMatrix, K::GeneralizedKroneckerProduct)

In-place collection of `K` in `C`. If possible, specialized routines are used to
speed up the computation. The fallback is an element-wise iteration. In this case,
this function might be slow.
"""
function collect!(C::AbstractMatrix, K::GeneralizedKroneckerProduct)
    size(C) == size(K) || throw(DimensionMismatch("`K` $(size(K)) cannot be collected in `C` $(size(C))"))
    @inbounds for I in CartesianIndices(K)
        C[I] = K[I]
    end
    return C
end

"""
    collect(K::GeneralizedKroneckerProduct)

Collects a lazy instance of the `GeneralizedKroneckerProduct` type into a dense,
native matrix. Falls back to the element-wise case when not specialized method
is defined.
"""
function collect(K::GeneralizedKroneckerProduct{T}) where {T}
    C = Matrix{T}(undef, size(K)...)
    collect!(C, K)
    return C
end

"""
    KroneckerProduct{T,TA<:AbstractMatrix, TB<:AbstractMatrix} <: AbstractKroneckerProduct{T}

Concrete Kronecker product between two matrices `A` and `B`.
"""
struct KroneckerProduct{T<:Any, TA<:AbstractMatrix, TB<:AbstractMatrix} <: AbstractKroneckerProduct{T}
    A::TA
    B::TB
    function KroneckerProduct(A::AbstractMatrix{T}, B::AbstractMatrix{V}) where {T, V}
        return new{promote_type(T, V), typeof(A), typeof(B)}(A, B)
    end
end

# conversion
Base.convert(::Type{T}, K::KroneckerProduct) where {T<:KroneckerProduct} = K isa T ? K : T(K)
function KroneckerProduct{T,TA,TB}(K::KroneckerProduct) where {T,TA,TB}
    A, B = getmatrices(K)
    KroneckerProduct(convert(TA, A), convert(TB, B))
end

"""
    kronecker(A::AbstractMatrix, B::AbstractMatrix)

Construct a Kronecker product object between two arrays. Does not evaluate the
Kronecker product explictly.
"""
kronecker(A::AbstractMatrix, B::AbstractMatrix) = KroneckerProduct(A, B)

# version that have a vector as input reshape to matrices
kronecker(A::AbstractVector, B::AbstractMatrix) = KroneckerProduct(reshape(A,:,1), B)
kronecker(A::AbstractMatrix, B::AbstractVector) = KroneckerProduct(A, reshape(B,:,1))
kronecker(A::AbstractVector, B::AbstractVector) = KroneckerProduct(reshape(A,:,1), reshape(B,:,1))

"""
    kronecker(A::AbstractMatrix, B::AbstractMatrix)

Higher-order Kronecker lazy kronecker product, e.g.
```
kronecker(A, B, C, D)
```
"""
kronecker(A::AbstractVecOrMat, B::AbstractVecOrMat...) = kronecker(A,
                                                            kronecker(B...))

"""
    ⊗(A::AbstractMatrix, B::AbstractMatrix)

Binary operator for `kronecker`, computes as Lazy Kronecker product. See
`kronecker` for documentation.
"""
⊗(A::AbstractVecOrMat, B::AbstractVecOrMat) = kronecker(A, B)
⊗(A::AbstractVecOrMat...) = kronecker(A...)

"""
    getindex(K::AbstractKroneckerProduct, i1::Integer, i2::Integer)

Computes and returns the (i,j)-th element of an `AbstractKroneckerProduct` K.
Uses recursion if `K` is of an order greater than two.
"""
function getindex(K::AbstractKroneckerProduct, i1::Integer, i2::Integer)
    A, B = getmatrices(K)
    k, l = size(B)
    return (A[cld(i1, k), cld(i2, l)]::eltype(A)) * (B[(i1 - 1) % k + 1, (i2 - 1) % l + 1]::eltype(B))
end

"""
    getallfactors(K::AbstractKroneckerProduct)

Obtain all factors in an `AbstractKroneckerProduct` object.
"""
getallfactors(K::AbstractKroneckerProduct) = (getallfactors(K.A)..., getallfactors(K.B)...)
getallfactors(K::AbstractMatrix) = (K,)

"""
    getmatrices(K::AbstractKroneckerProduct)

Obtain the two matrices of an `AbstractKroneckerProduct` object.
"""
getmatrices(K::AbstractKroneckerProduct) = (K.A, K.B)

"""
    getmatrices(A::AbstractArray)

Returns a matrix itself. Needed for recursion.
"""
getmatrices(A::AbstractArray) = (A,)

firstmatrix(K::AbstractKroneckerProduct) = first(getmatrices(K))
lastmatrix(K::AbstractKroneckerProduct) = last(getmatrices(K))

"""
    size(K::AbstractKroneckerProduct)

Returns a the size of an `AbstractKroneckerProduct` instance.
"""
function size(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    (m, n) = size(A)
    (k, l) = size(B)
    return m * k, n * l
end

"""
    size(K::GeneralizedKroneckerProduct)

Returns a the size of an `GeneralizedKroneckerProduct` instance.
"""
size(K::GeneralizedKroneckerProduct, dim::Integer) = size(K)[dim]

# CHECKS

"""
    issquare(A::AbstractMatrix)

Checks if an array is a square matrix.
"""
function issquare(A::AbstractMatrix)
    m, n = size(A)
    return m == n
end

"""
    issquare(A::Factorization)

Checks if a Factorization struct represents a square matrix.
"""
function issquare(A::Factorization)
    m, n = size(A)
    return m == n
end

"""
    issymmetric(K::AbstractKroneckerProduct)

Checks if a Kronecker product is symmetric.
"""
function LinearAlgebra.issymmetric(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    return issymmetric(A) && issymmetric(B)
end

"""
    ishermitian(K::AbstractKroneckerProduct)

Checks if a Kronecker product is Hermitian.
"""
function LinearAlgebra.ishermitian(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    return ishermitian(A) && ishermitian(B)
end

"""
    isposdef(K::AbstractKroneckerProduct)

Test whether a Kronecker product is positive definite (and Hermitian) by trying
to perform a Cholesky factorization of K.
"""
function LinearAlgebra.isposdef(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    return isposdef(A) && isposdef(B)
end

# SCALAR PROPERTIES

"""
    order(M::AbstractMatrix)

Returns the order of a matrix, i.e. how many matrices are involved in the
Kronecker product (default to 1 for general matrices).
"""
order(M::AbstractMatrix) = 1
order(M::AbstractKroneckerProduct) = order(M.A) + order(M.B)

"""
    det(K::AbstractKroneckerProduct)

Compute the determinant of a Kronecker product.
"""
function LinearAlgebra.det(K::AbstractKroneckerProduct{T}) where {T}
    checksquare(K)
    A, B = getmatrices(K)
    if issquare(A) && issquare(B)
        m = size(A)[1]
        n = size(B)[1]
        return det(A)^n * det(B)^m
    else
        return zero(T)
    end
end

"""
    logdet(K::AbstractKroneckerProduct)

Compute the logarithm of the determinant of a Kronecker product.
"""
function LinearAlgebra.logdet(K::AbstractKroneckerProduct{T}) where {T}
    checksquare(K)
    A, B = getmatrices(K)
    if issquare(A) && issquare(B)
        m = size(A)[1]
        n = size(B)[1]
        return n * logdet(A) + m * logdet(B)
    else
        return real(T)(-Inf)
    end
end

"""
    tr(K::AbstractKroneckerProduct)

Compute the trace of a Kronecker product.
"""
function LinearAlgebra.tr(K::AbstractKroneckerProduct)
    checksquare(K)
    A, B = getmatrices(K)
    if issquare(A) && issquare(B)
        return tr(A) * tr(B)
    else
        return sum(diag(K))
    end
end

"""
    inv(K::AbstractKroneckerProduct)

Compute the inverse of a Kronecker product.
"""
function inv(K::AbstractKroneckerProduct)
    checksquare(K)
    A, B = getmatrices(K)
    if issquare(A) && issquare(B)
        return KroneckerProduct(inv(A), inv(B))
    else
        throw(SingularException(1))
    end
end

"""
    pinv(K::AbstractKroneckerProduct)

Compute the Moore-Penrose pseudo-inverse of a Kronecker product.
"""
function LinearAlgebra.pinv(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    return KroneckerProduct(pinv(A), pinv(B))
end

"""
    adjoint(K::AbstractKroneckerProduct)

Compute the adjoint of a Kronecker product.
"""
function adjoint(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    return kronecker(A', B')
end

"""
    transpose(K::AbstractKroneckerProduct)

Compute the transpose of a Kronecker product.
"""
function Base.transpose(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    return kronecker(transpose(A), transpose(B))
end

function Base.conj(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    return kronecker(conj(A), conj(B))
end

function LinearAlgebra.diag(K::KroneckerProduct)
    A, B = getmatrices(K)
    if issquare(A) && issquare(B)
        return kron(diag(K.A), diag(K.B))
    end
    return K[diagind(K)]
end

# COLLECTING

#=
"""
    collect(K::AbstractKroneckerProduct)

Collects a lazy instance of the `AbstractKroneckerProduct` type into a full,
native matrix. Equivalent with `Matrix(K::AbstractKroneckerProduct)`.
"""
function collect(K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    return kron(A, B)
end
=#

_maybecollect(A::GeneralizedKroneckerProduct) = collect(A)
_maybecollect(A::AbstractArray) = A

# function for in-place Kronecker product
function _kron!(C::AbstractArray, A::AbstractArray, B::AbstractArray)
    m = first(LinearIndices(C)) - 1
    @inbounds for j = axes(A,2), l = axes(B,2), i = axes(A,1)
        Aij = A[i,j]
        for k = axes(B,1)
            C[m += 1] = Aij * B[k,l]
        end
    end
    return C
end

_kron!(C::AbstractArray, A::GeneralizedKroneckerProduct, B::AbstractArray) = _kron!(C, collect(A), B)
_kron!(C::AbstractArray, A::AbstractArray, B::GeneralizedKroneckerProduct) = _kron!(C, A, collect(B))
_kron!(C::AbstractArray, A::GeneralizedKroneckerProduct, B::GeneralizedKroneckerProduct) = _kron!(C, collect(A), collect(B))

@inline function _kron!(C::AbstractArray, As::Tuple{AbstractArray, Vararg{AbstractArray}}, Bs::Tuple{AbstractArray, Vararg{AbstractArray}}, f = identity)
    m = first(LinearIndices(C)) - 1
    A1 = first(As)
    B1 = first(Bs)
    for j = axes(A1,2), l = axes(B1,2), i = axes(A1,1)
        Aijs = map(A -> A[i,j], As)
        for k = 1:size(Bs[1],1)
            Bkls = map(B -> B[k,l], Bs)
            Aijs_times_Bkls = map(*, Aijs, Bkls)
            C[m += 1] = f(Aijs_times_Bkls...)
        end
    end
    return C
end

"""
    collect!(C::AbstractMatrix, K::AbstractKroneckerProduct)

In-place collection of `K` in `C` where `K` is an `AbstractKroneckerProduct`, i.e.,
`K = A ⊗ B`. This is equivalent to the broadcasted assignment `C .= K`.

    collect!(f, C::AbstractMatrix, K1::AbstractKroneckerProduct, Ks::AbstractKroneckerProduct...)

Evaluate `f.(K1, Ks...)` and assign it in-place to `C`. This is equivalent to the broadcasted
operation `C .= f.(K1, Ks...)`.
"""
function collect!(C::AbstractMatrix, K::AbstractKroneckerProduct)
    size(C) == size(K) || throw(DimensionMismatch("`K` $(size(K)) cannot be collected in `C` $(size(C))"))
    A, B = getmatrices(K)
    return _kron!(C, A, B)
end

@inline function collect!(f, C::AbstractMatrix, K1::AbstractKroneckerProduct, Ks::AbstractKroneckerProduct...)
    @noinline throwdm(K1sz, Csz) = throw(DimensionMismatch("`K` $K1sz cannot be collected in `C` $Csz"))
    size(C) == size(K1) || throwdm(size(K1), size(C))
    for K in Ks
        size(C) == size(K) || throwdm(size(K), size(C))
    end
    Ks_all = (K1, Ks...)
    As = map(x -> first(getmatrices(x)), Ks_all)
    Bs = map(x -> last(getmatrices(x)), Ks_all)
    return _kron!(C, As, Bs, f)
end

"""
    Matrix(K::GeneralizedKroneckerProduct)

Converts a `GeneralizedKroneckerProduct` instance to a Matrix type.
"""
Base.Matrix(K::GeneralizedKroneckerProduct) = collect(K)

function Base.:+(A::AbstractKroneckerProduct, B::StridedMatrix)
    T = promote_type(eltype(A), eltype(B))
    C = similar(Array{T}, size(A))
    C .= A
    C .+= B
    return C
end
Base.:+(A::StridedMatrix, B::AbstractKroneckerProduct) = B + A

function Base.:+(A::AbstractKroneckerProduct, B1::AbstractKroneckerProduct, Brest::AbstractKroneckerProduct...)
    Bs = (B1, Brest...)
    for B in Bs
        Base.promote_shape(A, B) # check size compatibility
    end
    # special methods to handle kronecker products with singleton matrices
    # if one matrix is common to all products, we only need to add the other matrix
    if all(x -> firstmatrix(x) === firstmatrix(A), Bs)
        K1 = kronecker(firstmatrix(A), +(lastmatrix(A), map(lastmatrix, Bs)...))
        return collect(K1)
    elseif all(x -> lastmatrix(x) === lastmatrix(A), Bs)
        K2 = kronecker(+(firstmatrix(A), map(firstmatrix, Bs)...), lastmatrix(A))
        return collect(K2)
    end
    # if the sizes of the component matrices are compatible, the operation may be
    # short-circuited
    sa = map(size, getmatrices(A))
    if all(x -> map(size, getmatrices(x)) == sa, Bs)
        return broadcast(+, A, Bs...)
    end
    # collect the arrrays before adding to avoid indexing into the kronecker products
    return +(collect(A), map(collect, Bs)...)
end

function Base.:-(A::AbstractKroneckerProduct, B::StridedMatrix)
    T = promote_type(eltype(A), eltype(B))
    C = similar(Array{T}, size(A))
    C .= A
    C .-= B
    return C
end
function Base.:-(A::StridedMatrix, B::AbstractKroneckerProduct)
    T = promote_type(eltype(A), eltype(B))
    C = similar(Array{T}, size(A))
    @. C = -B
    C .+= A
    return C
end

const KronProdDiagonal = KroneckerProduct{<:Any, <:Diagonal, <:Diagonal}
for T in [:Diagonal, :UniformScaling]
    @eval Base.:+(K::KronProdDiagonal, D::$T) = Diagonal(K) + D
    @eval Base.:+(D::$T, K::KronProdDiagonal) = D + Diagonal(K)
    @eval Base.:-(K::KronProdDiagonal, D::$T) = Diagonal(K) - D
    @eval Base.:-(D::$T, K::KronProdDiagonal) = D - Diagonal(K)
end

function Base.:-(A::AbstractKroneckerProduct, B::AbstractKroneckerProduct)
    # special methods to handle kronecker products with singleton matrices
    # if one matrix is common to all products, we only need to add the other matrix
    if firstmatrix(B) === firstmatrix(A)
        K1 = kronecker(firstmatrix(A), lastmatrix(A) - lastmatrix(B))
        return collect(K1)
    elseif lastmatrix(B) === lastmatrix(A)
        K2 = kronecker(firstmatrix(A) - firstmatrix(B), lastmatrix(A))
        return collect(K2)
    end
    # if the sizes of the component matrices are compatible, the operation may be
    # short-circuited
    if map(size, getmatrices(A)) == map(size, getmatrices(B))
        return A .- B
    end
    # collect the arrrays before subtracting to avoid indexing into the kronecker products
    return collect(A) - collect(B)
end

function Base.kron(K::AbstractKroneckerProduct, C::AbstractMatrix)
    A, B = getmatrices(K)
    return kron(kron(A, B), C)
end

function Base.kron(A::AbstractMatrix, K::AbstractKroneckerProduct)
    B, C = getmatrices(K)
    return kron(A, kron(B, C))
end

Base.kron(K1::AbstractKroneckerProduct,
            K2::AbstractKroneckerProduct) = kron(collect(K1), collect(K2))

# mixed-product property
function Base.:*(K1::AbstractKroneckerProduct, K2::AbstractKroneckerProduct)
    A, B = getmatrices(K1)
    C, D = getmatrices(K2)
    # check for size
    if size(A, 2) != size(C, 1)
        throw(DimensionMismatch("Mismatch between A and C in (A ⊗ B)(C ⊗ D)"))
    end
    if size(B, 2) != size(D, 1)
        throw(DimensionMismatch("Mismatch between B and D in (A ⊗ B)(C ⊗ D)"))
    end
    return (A * C) ⊗ (B * D)
end

"""
    lmul!(a::Number, K::AbstractKroneckerProduct)

Scale an `AbstractKroneckerProduct` `K` inplace by a factor `a` by rescaling the
**left** matrix.
"""
function LinearAlgebra.lmul!(a::Number, K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    lmul!(a, A)
end

"""
    rmul!(K::AbstractKroneckerProduct, a::Number)

Scale an `AbstractKroneckerProduct` `K` inplace by a factor `a` by rescaling the
**right** matrix.
"""
function LinearAlgebra.rmul!(K::AbstractKroneckerProduct, a::Number)
    A, B = getmatrices(K)
    rmul!(B, a)
end

function Base.:*(a::Number, K::AbstractKroneckerProduct)
    A, B = getmatrices(K)
    kronecker(a * A, B)
end

function Base.:*(K::AbstractKroneckerProduct, a::Number)
    A, B = getmatrices(K)
    kronecker(A, B * a)
end

function LinearAlgebra.power_by_squaring(K::KroneckerProduct, p::Integer)
    A, B = getmatrices(K)
    kronecker(A^p, B^p)
end

function LinearAlgebra.svdvals(K::KroneckerProduct)
    A, B = getmatrices(K)
    σA = svdvals(A)
    σB = svdvals(B)
    σ = sort!(kron(σA, σB), rev = true)
    n = minimum(size(K))
    if length(σ) < n
        append!(σ, zeros(eltype(σ), n - length(σ)))
    end
    return σ
end

# Broadcasting machinery

Base.copyto!(dest::AbstractMatrix, K::AbstractKroneckerProduct) = collect!(dest, K)

struct AbsKronProdStyle <: Broadcast.AbstractArrayStyle{2} end
AbsKronProdStyle(::Val{N}) where {N} = Broadcast.DefaultArrayStyle{N}()
AbsKronProdStyle(::Val{2}) = AbsKronProdStyle()

Base.BroadcastStyle(::Type{<:AbstractKroneckerProduct}) = AbsKronProdStyle()

function Base.similar(bc::Broadcast.Broadcasted{AbsKronProdStyle}, ::Type{T}) where {T}
    similar(Array{T}, axes(bc))
end

@inline function Base.copyto!(dest::AbstractArray, bc::Broadcast.Broadcasted{AbsKronProdStyle})
    @noinline throwdm(axdest, axsrc) =
        throw(DimensionMismatch("destination axes $axdest are not compatible with source axes $axsrc"))
    axes(dest) == axes(bc) || throwdm(axes(dest), axes(bc))
    # Some common cases that may be short-circuited
    # Case 1, example: A .= B
    if bc.args isa Tuple{AbstractKroneckerProduct}
        A = bc.args[1]
        collect!(bc.f, dest, A)
        return dest
    # Case 2, example: 2 .* B
    elseif bc.args isa Tuple{Number, AbstractKroneckerProduct}
        A = last(bc.args)
        n = first(bc.args)
        collect!(let n = n; x -> bc.f(n, x); end, dest, A)
        return dest
    # Case 3, example: B .* 2
    elseif bc.args isa Tuple{AbstractKroneckerProduct, Number}
        A = first(bc.args)
        n = last(bc.args)
        collect!(let n = n; x -> bc.f(x, n); end, dest, A)
        return dest
    end
    # An operation like K1 .+ K2 may be short-circuited if the component matrices
    # have the same size.
    bcf = Broadcast.flatten(bc)
    if all(x -> x isa AbstractKroneckerProduct, bcf.args)
        sz1 = map(size, getmatrices(bcf.args[1]))
        if all(x -> map(size, getmatrices(x)) == sz1, bcf.args[2:end])
            collect!(bcf.f, dest, bcf.args...)
            return dest
        end
    elseif all(x -> x isa Union{AbstractKroneckerProduct, StridedArray}, bcf.args)
        # Performance is better if the kronecker products are collected before the
        # broadcasted operation is performed, although this incurs allocations
        broadcast!(bcf.f, dest, map(_maybecollect, bcf.args)...)
        return dest
    end
    # The general case that indexes into each array.
    # This may be slow as indexing into an AbstractKroneckerProduct is expensive
    bc′ = Broadcast.preprocess(dest, bc)
    # Performance may vary depending on whether `@inbounds` is placed outside the
    # for loop or not. (cf. https://github.com/JuliaLang/julia/issues/38086)
    @inbounds @simd for I in eachindex(bc′)
        dest[I] = bc′[I]
    end
    return dest
end


#Demo 4: Multi-dimensional datasets (advanced usage)
#-------------------------------------------------------------------------------

using CMDimData
using CMDimData.MDDatasets
using CMDimData.EasyPlot
import Printf: @sprintf

#==Attributes
===============================================================================#
LBL_AXIS_TIME = "Time (s)"
LBL_AXIS_AMPLITUDE = "Amplitude"

#==Input data
===============================================================================#
tfund = 1e-9 #Fundamental
osr = 20 #(fundamental) oversampling ratio
nfund = 20 #cycles of the fundamental

#==Computations
===============================================================================#
t = DataF1(0:(tfund/osr):(nfund*tfund))
tmax = maximum(t)

#Generate parameter sweeps:
sweeplist = PSweep[
	PSweep("period", [1,2,3]*tfund)
	PSweep("slope", [-1,0,1]*(1/tmax))
	PSweep("offset", [-.4, 0, .9])
]

#Generate data:
lines = DataHR{DataF1}(sweeplist) #Create empty dataset
tones = DataHR{DataF1}(sweeplist) #Create empty dataset
for inds in subscripts(lines)
	(T, m, b) = coordinates(lines, inds)
	lines.elem[inds...] = DataF1(t.x, (x)->m*x+b)
	tones.elem[inds...] = sin(t*(2pi/T))
end

tones += lines #Create shifted dataset
#Filter on 2T harmonic:
tones_2T = getsubarray(tones, period=2tfund)
#Filter on 3T harmonic:
tones_3T = getsubarray(tones, period=3tfund)
#Filter on increasing slope (3rd index):
tones_incm = getsubarray(tones, :,3,:)

#==Generate plot
===============================================================================#
strns(T) = @sprintf("%.1f ns", T/1e-9) #Generate string with # of ns from supplied value

plot = cons(:plot, nstrips=4, legend=false, #Legend too busy with GracePlot
	ystrip1 = set(axislabel=LBL_AXIS_AMPLITUDE, striplabel="Tones"),
	ystrip2 = set(axislabel=LBL_AXIS_AMPLITUDE, striplabel="Tones ($(strns(2tfund)))"),
	ystrip3 = set(axislabel=LBL_AXIS_AMPLITUDE, striplabel="Tones ($(strns(3tfund)))"),
	ystrip4 = set(axislabel=LBL_AXIS_AMPLITUDE, striplabel="Tones (increasing slope)"),
	xaxis = set(label=LBL_AXIS_TIME),
)
push!(plot,
	cons(:wfrm, tones, strip=1),
	cons(:wfrm, tones_2T, strip=2),
	cons(:wfrm, tones_3T, strip=3),
	cons(:wfrm, tones_incm, strip=4),
)

pcoll = cons(:plot_collection, title="Mulit-Dataset Tests: Subarrays", ncolumns=1)
	push!(pcoll, plot)

#==Return pcoll to user (call evalfile(...))
===============================================================================#
pcoll


# ------------------------------------------------------------------
# Copyright (c) 2018, Júlio Hoffimann Mendes <juliohm@stanford.edu>
# Licensed under the ISC License. See LICENCE in the project root.
# ------------------------------------------------------------------

@userplot ExcessPlot

@recipe function f(p::ExcessPlot)
  # get user input
  xs = p.args[1]

  ks = 2:length(xs)
  ξs = meanexcess(xs, ks)

  seriestype --> :path
  xlabel --> "number of maxima"
  ylabel --> "extreme value index"

  ks, ξs
end


module MultiFieldFESpacesTests

using BlockArrays
using FillArrays
using Gridap.Arrays
using Gridap.Geometry
using Gridap.FESpaces
using Gridap.Fields
using Gridap.Integration
using Gridap.ReferenceFEs
using Gridap.CellData
using Test

using Gridap.MultiField
using Gridap.Arrays: BlockArrayCooMap

order = 2

domain = (0,1,0,1)
partition = (3,3)
model = CartesianDiscreteModel(domain,partition)

trian = get_triangulation(model)
degree = order
quad = CellQuadrature(trian,degree)

V = TestFESpace(model,ReferenceFE(lagrangian,Float64,order);conformity=:H1)
Q = TestFESpace(model,ReferenceFE(lagrangian,Float64,order-1),conformity=:L2)

U = TrialFESpace(V)
P = TrialFESpace(Q)

multi_field_style = ConsecutiveMultiFieldStyle()

Y = MultiFieldFESpace(Vector{Float64},[V,Q],multi_field_style)
X = MultiFieldFESpace(Vector{Float64},[U,P],multi_field_style)

@test num_free_dofs(X) == num_free_dofs(U) + num_free_dofs(P)
@test num_free_dofs(X) == num_free_dofs(Y)

dy = get_cell_shapefuns(Y)
dv, dq = dy

dx = get_cell_shapefuns_trial(X)
du, dp = dx

cellmat = integrate(dv*du,quad)
cellvec = integrate(dv*2,quad)
cellids = get_cell_to_bgcell(trian)
cellmatvec = pair_arrays(cellmat,cellvec)
@test isa(cellmat, LazyArray{<:Fill{<:BlockArrayCooMap}})
@test is_nonzero_block(cellmat[1],1,1)
@test is_zero_block(cellmat[1],1,2)
@test isa(cellvec, LazyArray{<:Fill{<:BlockArrayCooMap}})
@test is_nonzero_block(cellvec[1],1)
@test is_zero_block(cellvec[1],2)

matvecdata = (cellmatvec,cellids,cellids)
matdata = (cellmat,cellids,cellids)
vecdata = (cellvec,cellids)

free_values = rand(num_free_dofs(X))
xh = FEFunction(X,free_values)
test_fe_function(xh)
@test isa(xh,FEFunction)
uh, ph = xh
@test isa(uh,FEFunction)
@test isa(ph,FEFunction)

cell_isconstr = get_cell_isconstrained(X)
@test cell_isconstr == Fill(false,num_cells(model))

cell_constr = get_cell_constraints(X)
@test isa(cell_constr,LazyArray{<:Fill{<:BlockArrayCooMap}})

cell_dof_ids = get_cell_dof_ids(X)
@test isa(cell_dof_ids,LazyArray{<:Fill{<:BlockArrayCooMap}})

cf = CellField(X,get_cell_dof_ids(X))
@test isa(cf,MultiFieldCellField)

test_fe_space(X,matvecdata,matdata,vecdata)
test_fe_space(Y,matvecdata,matdata,vecdata)

#using Gridap.Visualization
#writevtk(trian,"trian";nsubcells=30,cellfields=["uh" => uh, "ph"=> ph])

f(x) = sin(4*pi*(x[1]-x[2]^2))+1
fh = interpolate([f,f],X)
fh = interpolate_everywhere([f,f],X)
fh = interpolate_dirichlet([f,f],X)

end # module


using PyPlot
using JLD
function plot_lines()
	X = load("../data/sens/solenoid/dJds_s3_K12.jld")
	rho = X["s"]
	dJds = X["dJds"]

	X = load("../data/obj_erg_avg/solenoid/r2_s3.jld")
    s3 = X["s3"]
    J = X["J"]
    fig, ax = subplots(1,1)
    ax.plot(s3, J, ".", ms=10.0)
    ax.xaxis.set_tick_params(labelsize=28)
    ax.yaxis.set_tick_params(labelsize=28)
    ax.set_xlabel(L"$s_2$",fontsize=28)
    ax.set_ylabel(L"$\langle J\rangle$",fontsize=28)
    ax.grid(true)

    eps = 1.5e-2

	X = load("../data/obj_erg_avg/solenoid/r2_s3_sens.jld")
	J = X["J"]
	s3 = X["s3"]

	n = size(dJds)[1]
	J_pts = reshape([J .- eps*dJds J .+ eps*dJds], n,
                    2)'
    s_pts = reshape([s3 .- eps s3 .+ eps], n,
                    2)'

    ax.plot(s_pts[:,1:n-1], J_pts[:,1:n-1], "k",lw=3.0)
    ax.plot(s_pts[:,n], J_pts[:,n], "k",lw=4.0,label=L"$d_s\langle J\rangle$ from S3")

    ax.legend(fontsize=28)

end


function plot_basic_state(y, dQdy, U, Ψ, Ly, file)

    p1 = plot(
        y,
        Ψ,
        linewidth = 3,
        label = "Ψ",
        xlabel = "y",
        title = "Streamfunction",
        xlims = (-Ly/2, Ly/2)
        )

    p2 = plot(
        y,
        U,
        linewidth = 3,
        label = "U",
        title = "Velocity",
        xlims = (-Ly/2, Ly/2)
        )

    p3 = plot(
        y,
        dQdy,
        linewidth = 3,
        label = "dQ/dy",
        xlabel = "y",
        title = "dQ/dy",
        xlims = (-Ly/2, Ly/2)
        )

    plt1 = plot(p1, p2, p3, layout = (3,1))
    savefig(plt1, file)

end

function plot_growth_rates(ks, σ, Nmodes, file)

    titles = ["1st mode", "2nd mode"]
    plt2 = plot()
    for cnt in 1:Nmodes
        plot!(plt2,
        ks,
        σ[cnt,:],
        linewidth = 3,
        label=titles[cnt],
        title="Growth Rates of QG Model",
        xlabel="k",
        xlims=(ks[1], ks[end]),
        xticks=(0:0.25:ks[end]),
        gridlinewidth=2.0
        )
    end

    savefig(plt2,file)

end

function plot_1D_streamfunction(k_index, k, Ny, y, σmodes, mode_number, file)

      Ly = params.Ly
      Lx = 2*pi/k;
      Nx = Ny;
      dx = Lx/Nx;
       x = collect(0:dx:Lx);
    X, Y = meshgrid(x,y)

    # Divide eigenvector matrix into specified spaces: u,v,eta.
    ψvec = σmodes[:, mode_number, k_index];

    plt = plot(
        y,
        real(ψvec),
        linewidth = 3,
        label = "real",
        xlabel = "y",
        title = "Streamfunction",
        xlims = (-Ly/2, Ly/2)
        )
    plot!(
        plt,
        y,
        linewidth = 3,
        imag(ψvec),
        label = "imag")
    savefig(plt, file)

end

function plot_2D_streamfunction(k_index, k, Ny, y, σmodes, mode_number, file)

    Ly = params.Ly
    Lx = 2*pi/k;
    Nx = Ny;
    dx = Lx/Nx;
     x = collect(0:dx:Lx+dx/2);
  X, Y = meshgrid(x,y[end:-1:1])

    ψvec = σmodes[1:Ny+1,        mode_number, k_index];

    ψ = repeat(real(ψvec),1,Nx+1).*cos.(k*X) - repeat(imag(ψvec),1,Nx+1).*sin.(k*X);

    kwargs = (
        xlabel = "x",
        ylabel = "y",
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :balance,
      colorbar = true,
          xlim = (    0, Lx),
         ylim =  (-Ly/2, Ly/2)
    )

    plt = contour(x, y[end:-1:1], ψ[:, end:-1:1], title="ψ"; kwargs...)

    savefig(plt, file)

end

function plot_2D_vorticity(k_index, k, Ny, y, Dy2, σmodes, mode_number, file)

    Ly = params.Ly
    Lx = 2*pi/k;
    Nx = Ny;
    dx = Lx/Nx;
     x = collect(0:dx:Lx+dx/2);
  X, Y = meshgrid(x,y[end:-1:1])

    ψvec = σmodes[1:Ny+1,        mode_number, k_index];

    ζvec = Dy2*ψvec - k*k*ψvec

    ζ = repeat(real(ζvec),1,Nx+1).*cos.(k*X) - repeat(imag(ζvec),1,Nx+1).*sin.(k*X);

    kwargs = (
        xlabel = "x",
        ylabel = "y",
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :balance,
      colorbar = true,
          xlim = (    0, Lx),
         ylim =  (-Ly/2, Ly/2)
    )

    plt = contour(x, y[end:-1:1], ζ[:, end:-1:1], title="ζ"; kwargs...)

    savefig(plt, file)

end


module QuadraturesTests

using Test
using Gridap

const D = 2
quad = TensorProductQuadrature(orders=(2,4))

coords_ref = [
  -0.57735 0.57735 -0.57735 0.57735 -0.57735 0.57735;
  -0.774597 -0.774597 0.0 0.0 0.774597 0.774597]

@test isa(coordinates(quad),Array{Point{D,Float64},1})
@test isa(weights(quad),Array{Float64,1})

coords = coordinates(quad)

coords_arr = reshape( reinterpret(Float64,coords), (D,length(coords)) )

@test isapprox(coords_ref,coords_arr,rtol=10-3)

weigs = weights(quad)

weigs_ref = [0.555556, 0.555556, 0.888889, 0.888889, 0.555556, 0.555556]

@test isapprox(weigs,weigs_ref,rtol=10-3)

quad = Quadrature((HEX_AXIS,HEX_AXIS,HEX_AXIS),order=2)

@test isa(quad,TensorProductQuadrature{3})

end # module QuadraturesTests


# Incidence matrix for element connections in the grid
function create_incidence_matrix(g::Grid)
    cell_containing_node = Dict{Int, Set{Int}}()
    for (cellid, cell) in enumerate(g.cells)
        for v in cell.nodes
            if !haskey(cell_containing_node, v)
                cell_containing_node[v] = Set{Int}()
            end
            push!(cell_containing_node[v], cellid)
        end
    end

    I, J, V = Int[], Int[], Bool[]
    for (_, cells) in cell_containing_node
        for cell1 in cells # All these cells have a neighboring node
            for cell2 in cells
                # if true # cell1 != cell2
                if cell1 != cell2
                    push!(I, cell1)
                    push!(J, cell2)
                    push!(V, true)
                end
            end
        end
    end

    incidence_matrix = sparse(I, J, V)
    return incidence_matrix
end

# Greedy algorithm for coloring a grid such that no two cells with the same node
# have the same color
function greedy_coloring(incidence_matrix, cells=1:size(incidence_matrix, 1))
    cell_colors = Dict{Int,Int}(i => 0 for i in cells) # Zero represents no color set yet
    occupied_colors = Set{Int}()
    final_colors = Vector{Int}[]
    total_colors = 0
    for cellid in cells
        empty!(occupied_colors)
        # loop over neighbors
        for r in nzrange(incidence_matrix, cellid)
            cell_neighbour = incidence_matrix.rowval[r]
            cell_neighbour in cells || continue # Only care about the subset given in cells
            color = cell_colors[cell_neighbour]
            if color != 0
                push!(occupied_colors, color)
            end
        end

        # occupied colors now contains all the colors we are not allowed to use
        free_color = 0
        for attempt_color in 1:total_colors
            if attempt_color ∉ occupied_colors
                free_color = attempt_color
                break
            end
        end
        if free_color == 0 # no free color found, need to bump max colors
            total_colors += 1
            free_color = total_colors
            push!(final_colors, Int[])
        end
        @assert free_color != 0
        cell_colors[cellid] = free_color
        push!(final_colors[free_color], cellid)
    end
    return final_colors
end

# See Appendix A in https://www.math.colostate.edu/%7Ebangerth/publications/2013-pattern.pdf
function workstream_coloring(incidence_matrix)
    ###################
    # 1. Partitioning #
    ###################
    zones = Set{Int}[]
    ## Zone 1: Just the first element
    push!(zones, Set{Int}(1))
    Z = 2
    Z0 = Set{Int}() # Dummy zone
    n_visited = 1
    ## Zone N: All elements with connection to elements in Zone N-1
    while true
        s = Set{Int}()
        # Loop over all elements in previous zone and add their neigbouring elements
        # unless they are in any of the previous 2 zones.
        for c in get(zones, Z-1, Z0)
            for r in nzrange(incidence_matrix, c)
                cell_neighbour = incidence_matrix.rowval[r]
                if !(cell_neighbour in get(zones, Z-2, Z0) || cell_neighbour in get(zones, Z-1, Z0))
                    push!(s, cell_neighbour)
                end
            end
        end
        push!(zones, s)
        n_visited += length(s)
        if n_visited >= size(incidence_matrix, 1)
            break
        end
        Z += 1
    end

    ###############
    # 2. Coloring #
    ###############
    # TODO: The reference uses DSATUR algorithm instead of greedy
    # TODO: Zones can be colorized in parallel
    zone_colors = Vector{Vector{Int}}[greedy_coloring(incidence_matrix, z) for z in zones]

    ################
    # 3. Gathering #
    ################
    Nodd,  Zodd  = findmax(x -> isodd(x)  ? length(zone_colors[x]) : typemin(Int), 1:length(zone_colors))
    Neven, Zeven = findmax(x -> iseven(x) ? length(zone_colors[x]) : typemin(Int), 1:length(zone_colors))
    N = Nodd + Neven
    final_colors = append!(zone_colors[Zodd], zone_colors[Zeven]) # Reuse these for output
    color_sizes = map(length, final_colors)
    used_for_zone = Set{Int}()
    for Z in 1:length(zone_colors)
        (Z == Zodd || Z == Zeven) && continue
        zone_color_vectors = zone_colors[Z]
        odd = isodd(Z)

        empty!(used_for_zone)

        for local_color in sortperm(zone_color_vectors; by=length, rev=true)
            cond = odd ? (x -> x <= Nodd) : (x -> x > Nodd)
            _, global_color = findmin(x -> (cond(x) && x ∉ used_for_zone) ? color_sizes[x] : typemax(Int), 1:N)
            push!(used_for_zone, global_color)
            append!(final_colors[global_color], zone_color_vectors[local_color])
            map!(length, color_sizes, final_colors)
        end
    end

    # Maybe nice to sort?
    foreach(sort!, final_colors)

    return final_colors
end

@enum ColoringAlgorithm GREEDY WORKSTREAM

"""
    create_coloring(g::Grid; alg::ColoringAlgorithm)

Create a coloring of the cells in grid `g` such that no neighboring cells
have the same color.

Returns a vector of vectors with cell indexes, e.g.:

```julia
ret = [
   [1, 3, 5, 10, ...], # cells for color 1
   [2, 4, 6, 12, ...], # cells for color 2
]
```

Two different algorithms are available, specified with the `alg` keyword argument:
 - `alg = Ferrite.WORKSTREAM` (default): Three step algorithm from
   [*WorkStream*](https://www.math.colostate.edu/%7Ebangerth/publications/2013-pattern.pdf)
   , albeit with a greedy coloring in the second step.
 - `alg = Ferrite.GREEDY`: greedy algorithm that works well for structured grid such as
   e.g. grids from `generate_grid`.

The resulting colors can be visualized using [`vtk_cell_data_colors`](@ref).

!!! note "Cell to color mapping"
    In a previous version of Ferrite this function returned a dictionary mapping
    cell ID to color numbers as the first argument. If you need this mapping you
    can create it using the following construct:
    ```julia
    colors = create_coloring(...)
    cell_colormap = Dict{Int,Int}(
        cellid => color for (color, cellids) in enumerate(final_colors) for cellid in cellids
    )
    ```
"""
function create_coloring(g::Grid; alg::ColoringAlgorithm=WORKSTREAM)
    incidence_matrix = create_incidence_matrix(g)
    if alg === WORKSTREAM
        return workstream_coloring(incidence_matrix)
    elseif alg === GREEDY
        return greedy_coloring(incidence_matrix)
    else
        error("impossible")
    end
end

"""
    vtk_cell_data_colors(vtkfile, cell_colors, name="coloring")

Write cell colors (see [`create_coloring`](@ref)) to a VTK file for visualization.
"""
function vtk_cell_data_colors(vtkfile, cell_colors::AbstractVector{<:AbstractVector{<:Integer}}, name="coloring")
    color_vector = zeros(sum(length, cell_colors))
    for (i, cells_color) in enumerate(cell_colors)
        for cell in cells_color
            color_vector[cell] = i
        end
    end
    vtk_cell_data(vtkfile, color_vector, name)
end


# EXAMPLE INPUT:
example = [
1721,
979,
366,
299,
675,
1456,
]

# Parse input
f = open("1.txt", "r")
lines = readlines(f)
numbers = map((x) -> parse(Int64, x), lines)
close(f)

function compute2(d)
    for i in 1:length(d)
        for j in i + 1:length(d)
            if d[i] + d[j] == 2020
                return d[i] * d[j]
            end
        end
    end
    return -1
end

function compute3(d)
    for i in 1:length(d)
        for j in i + 1:length(d)
            for k in j + 1:length(d)
                if d[i] + d[j] + d[k] == 2020
                    return d[i] * d[j] * d[k]
                end
            end

        end
    end
    return -1
end

println(compute2(example))
println(compute3(example))

println(compute2(numbers))
println(compute3(numbers))


module MNIST
using GZip

const mnisturl = "http://yann.lecun.com/exdb/mnist"
const xtrn_file = "train-images-idx3-ubyte.gz"
const ytrn_file = "train-labels-idx1-ubyte.gz"
const xtst_file = "t10k-images-idx3-ubyte.gz"
const ytst_file = "t10k-labels-idx1-ubyte.gz"

function wgetzcat(gz)
    #isfile(gz) || get("$mnisturl/$gz"; ostream=gz) # run(`wget $mnisturl/$gz`)
    isfile(gz) || run(`wget $mnisturl/$gz`)
    fh = GZip.open(gz)
    a = readbytes(fh)
    close(fh)
    a
end

function readlabels(gz)
    a=convert(Vector{Int}, wgetzcat(gz)[9:end])
    a[a.==0]=10
    full(sparse(a, 1:length(a), 1.0f0))
end

function readimages(gz)
    a=(wgetzcat(gz)[17:end] ./ 255.0f0)
    reshape(a, 28*28, div(length(a),(28*28)))
end

function loadmnist()
    global xtrn, ytrn, xtst, ytst
    xtrn = readimages(xtrn_file)
    ytrn = readlabels(ytrn_file)
    xtst = readimages(xtst_file)
    ytst = readlabels(ytst_file)
end

info("Loading MNIST...")
loadmnist()

end # module


#=
Grammar:
<exp>    ::= <term><exp'>
<exp'>   ::= +<term><exp'> | - <term><exp'> | ϵ
<term>   ::= <factor><term'>
<term'>  ::= *<factor><term'> | / <factor><term'> | ϵ
<factor> ::= (<exp>) |<ident>|<num>
<ident>  ::= A|B|...|Z
<num>    ::  Real, let's assume all Float64 numbers
=#

#=
   Recursive descent parser.
   We do LL(1). That is left-to-right leftmost derivation
=#

#=
 Let's construct one method for each non terminal.
 Let's assume that we have the method scan that updates our global token.
=#

global identLevel = 0
global current_token = ""

function print_and_ident(str::String)
  ident = ""
  for i in 1:identLevel
    ident *= "\t"
  end
  println("$ident $str")
end

function scan()
  res = read(stdin, Char)
  global current_token = string(res)
  return string(res)
end

function main()
  scan()
  res = 0
  res = exp(res)
  println("Result of res is: $res")
end

function exp(res)
  res = term(res)
  res = expPrime(res)
end

function term(res)
  res = factor(res)
  termPrime(res)
end

function expPrime(res)
  if current_token == "+"
    scan()
    res = res + term(res)
    expPrime(res)
  elseif current_token == "-"
    scan()
    res = res - term(res)
    expPrime(res)
  else #= ϵ =#
    return res
  end
end

function termPrime(res)
  if current_token == "*"
    scan()
    res = res * factor(res)
    termPrime(res)
  elseif current_token == "/"
    scan()
    res = res / factor(res)
    termPrime(res)
  else #= ϵ =#
    return res
  end
end

function factor(res)
  if current_token == "("
    scan()
    res = exp(res)
    if current_token != ")"
      throw("Right parenthesis expected")
    end
    scan()
    return res
  elseif nothing != match(r"[A-Z]+", current_token)
    res = parse(Float64, current_token)
    scan()
    return res
  elseif nothing != match(r"[0-9]+", current_token)
    res = parse(Float64, current_token)
    scan()
    return res
  else #= Faulty input =#
    throw("Unexpected token: $(current_token)")
  end
end

main()


function funk1(A)
    (a,b) = size(A)
    s = 0
    for i in 1:a
        for j in 1:b
            if rem(i+j,2) == 0
                s+=A[i,j]
            end
        end
    end
    return s
end

function funk2(A)
    (a,b) = size(A)
    s = 0
    br = 0
    for i in 1:a
        for j in 1:b
            if(rem(j,2)==1)
                if rem(A[i,j],2)==0
                    s+=A[i,j]
                    br+=1
                end
            end
        end
    end
    return s/br
end

function funk3(A)
     range10 = zeros(0)
     range100 = zeros(0)
     range1000 = zeros(0)
    n = size(A)[2]
    for i in 1:n
        if A[i]>=0 && A[i] <10
            push!(range10,A[i])
        elseif  A[i]>=10 && A[i]<100
            push!(range100,A[i])
        elseif A[i]>=100 && A[i]<1000
            push!(range1000,A[i])
        end
    end
    return range10, range100, range1000
end


using LinearAlgebra
using Polynomials

include("../tools/plotting.jl")
"""
Use Taylor's Theorem with Remainder to check
validity of computed Jacobian of f: R^n -> R^m.

More specifically, check that error of linear approximation decays like O(h^2)

last: smallest power of exp^(-x) we wish to compute in scheme

Requirements: derivative/jacobian must be in the right 2D shape, and not be a 1D list

Hessian is only for functions from R^n -> R^1
"""
function checkDerivative(f, df, x0, hessian = nothing, first = 3, last = 12, num = 10)
    if size(x0, 2) != 1
        x0 = reshape(x0, length(x0), 1) #reshape x0 into 2D column vector or 1D Array
    end
    f0 = f(x0)
    df0 = df(x0)
    n = length(x0)
    m = length(f0)
    @assert size(f0, 2)==1 #no reason for f0 to be 2D
    #println(size(df0))
    #println("m: ", m)
    #println("n: ", n)
    #@assert size(df0) == (m, n) ||  size(df0) == (m, ) #keep 1D list as possibility
    #println("size df0, 1: ", size(df0, 1))
    #println("size df0, 2: ", size(df0, 2))
    try
        @assert size(df0, 1) == m && size(df0, 2) == n
    catch e
        #@info "df0", df0
        #@info "m", m
        #@info "n", n
    end

    try
        df0 = reshape(df0, m, n) #reshapes to m x n if non-scalar
    catch e
    end

    if hessian!=nothing
        @assert m == 1 #we can only compute the Hessian for functions to R^1
        d2f0 = hessian(x0)
        #println(size(d2f0))
        @assert size(d2f0, 1) == n && size(d2f0, 1) ==n
    end
    #what if x0 or f0 are constants?
    dx = rand(size(x0, 1), size(x0, 2))
    if length(f0) >1
        f0 = reshape(f0, length(f0), 1) #reshape into col vector
    end
    h = collect(first:(last-first)/num:last)
    for i=1:length(h)
        h[i] = exp(-h[i])
    end
    A = zeros(length(h))
    for i = 1:length(h)
        fi = f(x0 .+ h[i]*dx) #dx has the same shape x0, so we can safely do .+, which allows us to add scalar and array
                               #if dx did not have same shape as x0, then .+ could have unintended broadcasting effects
        @assert size(fi, 1) == size(f0, 1) && size(fi, 2) == size(f0, 2) #this lets us do fi .- f0 safely
        if false#debug
            println("x0: ", x0)
            println("dx: ", dx)
            println("fi: ", fi)
            println("f0: ", f0)
            println("df0: ", df0)
            println("increment", h[i]*dx)
            println("h[i]: ", h[i])
        end
            if hessian!=nothing
                inc = h[i] .* dx
                expr1 = fi .- f0
                expr2 = df0 * inc .+ 0.5* inc' * d2f0 * inc
                #println("size expr1: ", size(expr1))
                #println("size expr2: ", size(expr2))
                @assert size(expr1, 1) == size(expr2, 1) && size(expr1, 2) == size(expr2, 2)
                A[i] = norm(expr1 .- expr2)
                #A[i] = norm((fi - f0) - df0' * inc .- 0.5* inc' * d2f0 * inc)
            else
                expr1 = fi .- f0
                expr2 = df0 * (h[i] .* dx)
                #println("size expr1: ", size(expr1))
                #println("size expr2: ", size(expr2))
                @assert size(expr1, 1) == size(expr2, 1) && size(expr1, 2)==size(expr2, 2)
                A[i] = norm(expr1 .- expr2)
                #catch DimensionMismatch #we catch the case when f: R^1 -> R^n, in which case  df0'*dx will yield an error
                #    println("Warning: finitedifference dimension mismatch")
                #    @assert size(expr1) == size(expr2)
                #    A[i] = norm(expr1 - expr2)
                #end
        end
    end
    r1 = log.(h)
    r2 = log.(A)
    plt = Plots.plot(r1, r2, title = "Finite Difference Derivative Checker", xlabel = "log of h", ylabel = "log of error")#, fontfamily = font(48, "Courier") , reuse = false)
    return (r1, r2, plt, fit(r1, r2, 1))
end

"""
df is an in-place function (has no return value)
"""
function checkDerivative_in_place(f, df, x0, store, hessian = nothing, first = 3, last = 12, num = 10)
    if size(x0, 2) != 1
        x0 = reshape(x0, length(x0), 1) #reshape x0 into 2D column vector or 1D Array
    end
    f0 = f(x0)
    #df0 = df(x0)
    df(store, x0) #result in store
    n = length(x0)
    m = length(f0)
    @assert size(f0, 2)==1 #no reason for f0 to be 2D
    #println(size(df0))
    #println("m: ", m)
    #println("n: ", n)
    #@assert size(df0) == (m, n) ||  size(df0) == (m, ) #keep 1D list as possibility
    #println("size df0, 1: ", size(df0, 1))
    #println("size df0, 2: ", size(df0, 2))
    try
        @assert size(store, 1) == m && size(store, 2) == n
    catch e
        #@info "df0", store
        #@info "m", m
        #@info "n", n
        try
            store = reshape(store, m, n) # force reshape to m x n if non-scalar
        catch e
        end
    end
    if hessian!=nothing
        @assert m == 1 #we can only compute the Hessian for functions to R^1
        d2f0 = hessian(x0)
        #println(size(d2f0))
        @assert size(d2f0, 1) == n && size(d2f0, 1) ==n
    end
    #what if x0 or f0 are constants?
    dx = rand(size(x0, 1), size(x0, 2))
    if length(f0) >1
        f0 = reshape(f0, length(f0), 1) #reshape into col vector
    end
    h = collect(first:(last-first)/num:last)
    for i=1:length(h)
        h[i] = exp(-h[i])
    end
    A = zeros(length(h))
    for i = 1:length(h)
        fi = f(x0 .+ h[i]*dx) #dx has the same shape x0, so we can safely do .+, which allows us to add scalar and array
                               #if dx did not have same shape as x0, then .+ could have unintended broadcasting effects
        @assert size(fi, 1) == size(f0, 1) && size(fi, 2) == size(f0, 2) #this lets us do fi .- f0 safely
        if false#debug
            println("x0: ", x0)
            println("dx: ", dx)
            println("fi: ", fi)
            println("f0: ", f0)
            println("df0: ", store)
            println("increment", h[i]*dx)
            println("h[i]: ", h[i])
        end
            if hessian!=nothing
                inc = h[i] .* dx
                expr1 = fi .- f0
                expr2 = store * inc .+ 0.5* inc' * d2f0 * inc
                #println("size expr1: ", size(expr1))
                #println("size expr2: ", size(expr2))
                @assert size(expr1, 1) == size(expr2, 1) && size(expr1, 2) == size(expr2, 2)
                A[i] = norm(expr1 .- expr2)
                #A[i] = norm((fi - f0) - df0' * inc .- 0.5* inc' * d2f0 * inc)
            else
                expr1 = fi .- f0
                expr2 = store * (h[i] .* dx)
                #println("size expr1: ", size(expr1))
                #println("size expr2: ", size(expr2))
                @assert size(expr1, 1) == size(expr2, 1) && size(expr1, 2)==size(expr2, 2)
                A[i] = norm(expr1 .- expr2)
                #catch DimensionMismatch #we catch the case when f: R^1 -> R^n, in which case  df0'*dx will yield an error
                #    println("Warning: finitedifference dimension mismatch")
                #    @assert size(expr1) == size(expr2)
                #    A[i] = norm(expr1 - expr2)
                #end
        end
    end
    r1 = log.(h)
    r2 = log.(A)
    plt = Plots.plot(r1, r2, title = "Finite Difference Derivative Checker", xlabel = "log of h", ylabel = "log of error")#, fontfamily = font(48, "Courier") , reuse = false)
    return (r1, r2, plt, fit(r1, r2, 1))
end


function allornothing(network::AbstractNetwork,r,trips::Vector,costs::Vector)
    #nlinks = numlinks(network)
    zoneids = id.(zones(network))
    flows = zero(costs)

    _, parentvec = dijkstra(network, r, costs)
    for j in zoneids
        t = trips[j]
        if isapprox(t, 0.) | (r == j)
            continue
        end

        # Finding the shortest path
        p = pathto(parentvec, j)

        # Assigning flows
        for l in p
            flows[idx(network, l)] += t
        end

    end
    return flows
end

function allornothing(network::AbstractNetwork,
                      trips::Matrix{U},
                      costs::Vector{T};
                      basedon::Symbol = :link) where {T<:Real, U<:Real}
    # Initialization
    nnodes, nzones = numnodes(network), numzones(network)
    nlinks = length(costs)
    zonelist = zones(network)
    zoneids = id.(zonelist)

    flows = (basedon == :link) ? zeros(nlinks) : (basedon == :origin) ? zeros(nlinks, nzones) : error("Invalid input!")

    for (i,z) in enumerate(zoneids)
        if basedon == :link
            flows += allornothing(network, z, trips[z,:], costs)
        else
            flows[:,i] += allornothing(network, z, trips[z,:], costs)
        end
    end

    return flows
end

allornothing(network::AbstractNetwork, trips::AbstractMatrix{T}, costfn::CostFunction; basedon::Symbol = :link) where {T<:Real} = allornothing(network, trips, costfn(zeros(numlinks(network))).costs; basedon=basedon)


print("Last updated: $(read(`date`, String))")

# Load environment
import Pkg; Pkg.activate("../../../")

# Import Libraries
using Turing
using Turing: Variational
using Turing.RandomMeasures: DirichletProcess, StickBreakingProcess, stickbreak
using Distributions
using JSON3
using PyPlot
using StatsFuns
import Random
using BenchmarkTools
using Flux

# DP GMM model under stick-breaking construction
@model dp_gmm_sb(y, K) = begin
    nobs = length(y)

    mu ~ filldist(Normal(0, 3), K)
    sig ~ filldist(Gamma(1, 1/10), K)  # mean = 0.1

    alpha ~ Gamma(1, 1/10)  # mean = 0.1
    crm = DirichletProcess(alpha)
    v ~ filldist(StickBreakingProcess(crm), K - 1)
    eta = stickbreak(v)

    log_target = logsumexp(normlogpdf.(mu', sig', y) .+ log.(eta)', dims=2)
    Turing.acclogp!(_varinfo, sum(log_target))
end;

# Directory where all simulation data are stored.
data_dir = joinpath(@__DIR__, "../data")
path_to_data = joinpath(data_dir, "gmm-data-n200.json")

# Load data in JSON format.
data = let
    x = open(f -> read(f, String), path_to_data)
    JSON3.read(x, Dict{Symbol, Vector{Any}})
end

# Convert data to vector of floats
y = Float64.(data[:y]);

# Visualize data
plt.hist(y, bins=50, density=true)
plt.xlabel("y")
plt.ylabel("density")
plt.title("Histogram of data");

# Fit DP-SB-GMM with ADVI

# Set random seed for reproducibility
Random.seed!(10);

# Compile time approx: 1s
# Run time approx: 8s

# Define mean approximation for posterior
m = dp_gmm_sb(y, 10)
q0 = Variational.meanfield(m)  # initialize variational distribution (optional)
advi = ADVI(1, 2000)  # num_elbo_samples, max_iters
@time q = vi(m, advi, q0, optimizer=Flux.ADAM(1e-2));

# Function for generating samples from approximate posterior
nsamples = 1000
qsamples = rand(q, nsamples)
_, sym2range = Variational.bijector(m; sym_to_ranges = Val(true));
extract(sym) = qsamples[collect(sym2range[sym][1]), :]

# Extract eta
vpost = extract(:v)
etapost = hcat([stickbreak(vpost[:, col]) for col in 1:size(vpost, 2)]...)';

plt.figure(figsize=(7, 4))

plt.subplot(2, 2, 1)
plt.boxplot(etapost, whis=[2.5, 97.5], showmeans=true, showfliers=false);
foreach(line -> plt.axhline(line, ls=":"), data[:w])
plt.xlabel("mixture component")
plt.ylabel("η")

plt.subplot(2, 2, 2)
plt.boxplot(extract(:mu)', whis=[2.5, 97.5], showmeans=true, showfliers=false);
foreach(line -> plt.axhline(line, ls=":"), data[:mu])
plt.xlabel("mixture component")
plt.ylabel("μ")

plt.subplot(2, 2, 3)
plt.boxplot(extract(:sig)', whis=[2.5, 97.5], showmeans=true, showfliers=false);
foreach(line -> plt.axhline(line, ls=":"), data[:sig])
plt.xlabel("mixture component")
plt.ylabel("σ")

plt.subplot(2, 2, 4)
plt.hist(extract(:alpha)', density=true, bins=30)
plt.xlabel("α")
plt.ylabel("density")

plt.tight_layout();

function extract(chain, sym; burn=0, is_scalar=true)
    if is_scalar
        tail = chain[sym].data[(burn + 1):end, 1]
    else
        tail = group(chain, sym).value.data[(burn + 1):end, :, 1]
        # NOTE: `chain[:sym]` will not work in the future for parameter
        # vectors. Use `group` instead.
    end
    return tail
end

# Fit DP-SB-GMM with HMC

# Set random seed for reproducibility
Random.seed!(0);

# Compile time approx. 32s.
# Run time approx. 70s.

@time hmc_chain = begin
    burn = 500  # NOTE: The burn in is also returned. Can't be discarded.
    n_samples = 500
    iterations = burn + n_samples
    n_components = 10
    stepsize = 0.01
    nleapfrog = floor(Int, 1 / stepsize)

    sample(dp_gmm_sb(y, n_components),
           HMC(stepsize, nleapfrog),
           iterations)
end
;

# Fit DP-SB-GMM with NUTS

# Set random seed for reproducibility
Random.seed!(0);

# NUTS
# Compile time approx. 11s
# Run time approx. 244s
# Slower, but works a little better.
@time nuts_chain = begin
    n_components = 10
    n_samples = 500
    nadapt = 500
    iterations = n_samples + nadapt
    burn = 0  # For compatibility with HMC below.
    target_accept_ratio = 0.8

    sample(dp_gmm_sb(y, n_components),
           NUTS(nadapt, target_accept_ratio, max_depth=10),
           iterations);
end
;

function plot_param_post(param, param_name, param_full_name; figsize=(11, 4), truth=nothing)
    plt.figure(figsize=figsize)

    plt.subplot(1, 2, 1)
    plt.boxplot(param, whis=[2.5, 97.5], showmeans=true, showfliers=false)
    plt.xlabel("mixture components")
    plt.ylabel(param_full_name)
    plt.title("95% Credible Intervals for $(param_full_name)")

    if truth != nothing
        for line in truth
            plt.axhline(line, ls=":")
        end
    end

    plt.subplot(1, 2, 2)
    plt.plot(param)
    plt.xlabel("iterations")
    plt.ylabel(param_full_name)
    plt.title("Trace plot of $(param_full_name)");

    plt.tight_layout();
end

function plot_all_params(param; burn, figsize=(11, 4))
    vpost = extract(param, :v, burn=burn, is_scalar=false);
    mupost = extract(param, :mu, burn=burn, is_scalar=false);
    sigpost = extract(param, :sig, burn=burn, is_scalar=false);
    etapost = hcat([stickbreak(vpost[row, :]) for row in 1:size(vpost, 1)]...)';

    plt.figure(figsize=figsize)
    plt.subplot(1, 2, 1)

    # Loglikelihood can be extracted after model fitting using string macro.
    # See: https://turing.ml/dev/docs/using-turing/guide#querying-probabilities-from-model-or-chain

    # FIXME: not working on master branch.
    # loglike = logprob"y=y, K=n_components | model=dp_gmm_sb, chain=param"
    # plt.plot(loglike[(burn+1):end])
    # plt.xlabel("iteration (post-burn)")
    # plt.ylabel("Log likelihood")

    plt.subplot(1, 2, 2)
    plt.hist(extract(param, :alpha, burn=burn), density=true, bins=30)
    plt.xlabel("α")
    plt.ylabel("density")
    plt.title("Histogram of mass parameter α");
    plt.tight_layout();

    plot_param_post(etapost, :eta, "mixture weights (η)", truth=data[:w], figsize=figsize);
    plot_param_post(mupost, :mu, "mixture means (μ)", truth=data[:mu], figsize=figsize);
    plot_param_post(sigpost, :sigma, "mixture scales (σ)", truth=data[:sig], figsize=figsize);
end

plot_all_params(hmc_chain, burn=500, figsize=(8, 3));

plot_all_params(nuts_chain, burn=0, figsize=(8, 3));


"""
    equiSpacedPosition(x, knotsIn)
Returns the bin in which x is contained for the bins defined by the
limits in knotsIn. By default if the value is in the limit it is pPosition
at the right bin except for the las limit which is positioned to the left.
#Examples
```julia julia-repl
julia> equiSpacedPosition(2, [1,2,3,4,5,6,7,8])
2
```
"""
function equiSpacedPosition(x, knotsIn)

    @assert all(diff(knotsIn).>=0)    # bin limits are monotonically increaseing

    nbins = length(knotsIn) - 1
    lsubd = knotsIn[2] - knotsIn[1]

    proximityTol = 1e-4
    xinDist = abs(knotsIn[1]-x)
    xfinDist = abs(x - knotsIn[end])

    if x <= knotsIn[1]

        xinDist/lsubd < proximityTol ? s = 1 : throw(DomainError(x, "The value is out of the domain (left)"))

    elseif x >= knotsIn[end]

        xfinDist/lsubd < proximityTol ? s = nbins : throw(DomainError(x, "The value is out of the domain (right)"))

    else
        s = floor(Int,(x - knotsIn[1])/lsubd) + 1

        s > nbins ? s = nbins : nothing
    end
    return s
end

"""
    evalBSplineBasis(x, knots)
Builds the evaluation of the cubic BSpline basis corresponding to
the `knots` at `x`.
#Examples
```julia julia-repl
julia> evalBSplineBasis(10.,buildKnots(1.,10., 12))
12-element Vector{Float64}:
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.0
 0.16666666666666663
 0.6666666666666666
 0.16666666666666666
```
"""
function evalBSplineBasis(x, knots)
    transfMatrix = (1. / 6.) * [ -1. 3. -3. 1.;
                                  3. -6. 3. 0.;
                                 -3. 0. 3. 0.;
                                  1. 4. 1. 0.]
    nvertices = length(knots)
    lsubd = knots[2] - knots[1]
    knotsIn = knots[2:end-1] # the internal vertices are all except extrema

    bsplineBasis = zeros(promote_type(typeof(x), eltype(knots)), nvertices)

    s = equiSpacedPosition(x, knotsIn)

    ξ = (x - knotsIn[s])/lsubd
    ξvlocal = [ξ^3, ξ^2, ξ, one(promote_type(typeof(x), eltype(knots)))]

    bsplineBasis[s:s+3] = ξvlocal'*transfMatrix

    return bsplineBasis
end

"""
    evalBSplineBasis2dim(x, y, knotsx, knotsy)
Builds the evaluation of the dyadic product of 2 cubic BSplines
basis, the first corresponding to `knotsx` at `x` and the second
corresponding to `knotsy`at `y`. The output is already vectorized
instead of being a matrix.
#Examples
```julia julia-repl
julia> evalBSplineBasis2dim(3.,2.,buildKnots(1.,10., 12),buildKnots(0.,5.,14))
168-element Vector{Float64}:
 0.0
 0.0
 0.0
 0.0
 ⋮
 0.0
 0.0
 0.0
```
"""
function evalBSplineBasis2dim(x, y, knotsx, knotsy)

    bsplineBasisx = evalBSplineBasis(x, knotsx)
    bsplineBasisy = evalBSplineBasis(y, knotsy)
    return vec(bsplineBasisx*bsplineBasisy')
end

"""
    evalBSplineBasisD1(x, knots)
Builds the evaluation of the derivative of the cubic BSpline basis corresponding to
the `knots` at `x`.
#Examples
```julia julia-repl
julia> evalBSplineBasisD1(10.,buildKnots(1.,10., 12))
12-element Vector{Float64}:
  0.0
  0.0
  0.0
  0.0
  ⋮
 -0.5
  0.0
  0.5
```
"""
function evalBSplineBasisD1(x, knots)

    transfMatrix = (1. / 6.) * [ -1. 3. -3. 1.;
                                   3. -6. 3. 0.;
                                   -3. 0. 3. 0.;
                                    1. 4. 1. 0.]
    nvertices = length(knots)
    lsubd = knots[2] - knots[1]
    knotsIn = knots[2:end-1] # the internal vertices are all except extrema
    bsplineBasisD1 = zeros(typeof(x), nvertices)

    s = equiSpacedPosition(x, knotsIn)

    ξ = (x - knotsIn[s])/lsubd
    ξvlocalD1 = [3. * ξ^2, 2. * ξ, 1., 0.]

    bsplineBasisD1[s:s+3] = ξvlocalD1'*transfMatrix*(1/lsubd)
    return bsplineBasisD1
end

function evalBSplineBasis2dimD1x(x, y, knotsx, knotsy)
    dbsplineBasisx = evalBSplineBasisD1(x, knotsx)
    bsplineBasisy = evalBSplineBasis(y, knotsy)
    return vec(dbsplineBasisx*bsplineBasisy')
end

function evalBSplineBasis2dimD1y(x, y, knotsx, knotsy)
    bsplineBasisx = evalBSplineBasis(x, knotsx)
    dbsplineBasisy = evalBSplineBasisD1(y, knotsy)
    return vec(bsplineBasisx*dbsplineBasisy')
end

"""
    evalBSpline(x, knots, vertices)
Evaluation of the cubic BSpline at `x` for vertices values `vertices`
and the BSpline basis defined by `evalBSplineBasis(x, knots)`.
#Examples
```julia julia-repl
julia> evalBSpline(9.,buildKnots(1.,10., 12), ones(12))
0.9999999999999999
```
"""
function evalBSpline(x, knots, vertices)
    return evalBSplineBasis(x, knots)'*vertices
end

evalBSpline(x, bs::BSpline) = evalBSpline(x, bs.knots, bs.vertices)

"""
    evalBSpline2dim(x, y, knotsx, knotsy, vertices)
Evaluation of the cubic bidimensional BSpline with knots `knotsx`, `knotsy`
at `(x,y)` for vertices `vertices`.

#Examples
```julia julia-repl
julia> knotsx = knotsy = 0:0.1:10;
julia> vertices = rand(length(knotsx)*length(knotsy ));
julia> evalBSpline2dim(2., 3., knotsx, knotsy, vertices)
0.3800038047730809
```
"""
function evalBSpline2dim(x, y, knotsx, knotsy, vertices)
    return evalBSplineBasis2dim(x, y, knotsx, knotsy)'*vertices
end

evalBSpline2dim(x, y, bs2dim::BSpline2dim) = evalBSpline2dim(x, y, bs2dim.knotsx, bs2dim.knotsy, bs2dim.vertices)

"""
    evalBSplineD1(x, knots, vertices)
Evaluating the derivative ofthe cubic BSpline at `x` for vertices
values `vertices` and the BSpline basis derivative defined by
`evalBSplineBasisD1(x, knots)`.
#Examples
```julia julia-repl
julia> evalBSplineD1(9.,buildKnots(1.,10., 12), ones(12))
0.0
```
"""
function evalBSplineD1(x, knots, vertices)
    return evalBSplineBasisD1(x, knots)'*vertices
end

evalBSplineD1(x, bs::BSpline) = evalBSplineD1(x, bs.knots, bs.vertices)

"""
    evalBSpline2dimD1x(x, y, knotsx, knotsy, vertices)
Evaluation of the derivative with respect to `x`of the
cubic bidimensional BSpline with knots `knotsx`, `knotsy`
at `(x,y)` for vertices `vertices`.

#Examples
```julia julia-repl
julia> knotsx = knotsy = 0:0.1:10;
julia> vertices = rand(length(knotsx)*length(knotsy ));
julia> evalBSpline2dimD1x(2., 3., knotsx, knotsy, vertices)
-2.0863387254122316
```
"""
function evalBSpline2dimD1x(x, y, knotsx, knotsy, vertices)
    return evalBSplineBasis2dimD1x(x, y, knotsx, knotsy)'*vertices
end

evalBSpline2dimD1x(x, y, bs2dim::BSpline2dim) = evalBSpline2dimD1x(x, y, bs2dim.knotsx, bs2dim.knotsy, bs2dim.vertices)

"""
    evalBSpline2dimD1y(x, y, knotsx, knotsy, vertices)
Evaluation of the derivative with respect to `y`of the
cubic bidimensional BSpline with knots `knotsx`, `knotsy`
at `(x,y)` for vertices `vertices`.

#Examples
```julia julia-repl
julia> knotsx = knotsy = 0:0.1:10;
julia> vertices = rand(length(knotsx)*length(knotsy ));
julia> evalBSpline2dimD1y(2., 3., knotsx, knotsy, vertices)
2.345106484778582
```
"""
function evalBSpline2dimD1y(x, y, knotsx, knotsy, vertices)
    return evalBSplineBasis2dimD1y(x, y, knotsx, knotsy)'*vertices
end

evalBSpline2dimD1y(x, y, bs2dim::BSpline2dim) = evalBSpline2dimD1y(x, y, bs2dim.knotsx, bs2dim.knotsy, bs2dim.vertices)

"""
    nVerticesNeededLeftRight(knots, xin, xend)
Given the knots vector `knots` and the limits `xin` and `xend`,
determines the number of vertices that knots needs  on the left
and on the right needed to reach the values `xin` and `xend`.

#Examples
```julia julia-repl
julia> knots = 0.:1.:10.;
julia> nVerticesNeededLeftRight(knots, -2, 14)
(3, 5)
```
"""
function nVerticesNeededLeftRight(knots, xin, xend)
    knotsInmin = knots[2]
    knotsInmax = knots[end-1]

    lsubd = knots[2] - knots[1]

    intervLeft = (knotsInmin - xin)/lsubd
    intervRight = (xend - knotsInmax)/lsubd

    intervLeft > 0 ? nKnotsLeft = Int64(ceil(intervLeft)) : nKnotsLeft = 0
    intervRight > 0 ? nKnotsRight = Int64(ceil(intervRight)) : nKnotsRight = 0
    return nKnotsLeft, nKnotsRight
end

"""
    extendedKnots(knots, nKnotsLeft, nKnotsRight)
Given the knots vector `knots` and `nKnotsLeft` and `nKnotsRight`
extends the `knots` on the left `nKnotsLeft` and on the right `nKnotsRight`.

#Examples
```julia julia-repl
julia> knots = 0.:1.:10.;
julia> extendedKnots(knots, 2, 2)
15-element Vector{Float64}:
 -2.0
 -1.0
  ⋮
 11.0
 12.0
```
"""
function extendedKnots(knots, nKnotsLeft, nKnotsRight)

    lsubd = knots[2] - knots[1]

    addKnotsLeft = range(knots[1]-lsubd*nKnotsLeft, knots[1]-lsubd, length = nKnotsLeft)
    addKnotsRight = range(knots[end]+lsubd, knots[end]+lsubd*nKnotsRight, length = nKnotsRight)

    knotsNew = vcat(addKnotsLeft, knots, addKnotsRight)

    return knotsNew
end

"""
    expandBSpline(knots, vertices, xin, xend)
Given the knots and the vertices of a given BSpline with a certain range of validity,
it returns a new BSpline which extends the range of validity of the previous BSpline
betwenn `xin` and `xend`.
#Examples
```julia julia-repl
julia> bs2 = expandBSpline(bs1, 0., 5.);
julia> bs2.knots
20-element Vector{Float64}:
 -0.6666666666666663
 -0.33333333333333304
  2.220446049250313e-16
  ⋮
  5.333333333333333
  5.666666666666666
julia> bs2.vertices
20-element Vector{Float64}:
  -1.8150147193435557
  -1.2594147415453065
  -0.7038147637472744
   ⋮
  27.07391859179092
  29.851651947327774
```
"""
function expandBSpline(knots, vertices,
                       xin, xend)

    nvertices = length(knots)

    nKnotsLeft, nKnotsRight = nVerticesNeededLeftRight(knots, xin, xend)
    knotsNew = extendedKnots(knots, nKnotsLeft, nKnotsRight)
    nverticesNew = length(knotsNew)

    A = zeros(nverticesNew, nverticesNew)

    D2 = D2matrix(nverticesNew)
    Ω2 = Ω2matrix(nverticesNew)
    W = Wmatrix(nverticesNew)

    A[nKnotsLeft+1 : (nKnotsLeft + nvertices), nKnotsLeft+1 : (nKnotsLeft + nvertices)] =Float64.(I(nvertices))

    for i = nKnotsLeft+1:(nKnotsLeft + nvertices)
        W[i,i] = 1e4
    end

    Ω2[diagind(Ω2)] .= 1.

    Asys = A'*W*A + D2'*Ω2*D2
    bsys = A'*W*vcat(zeros(nKnotsLeft), vertices, zeros(nKnotsRight))

    verticesNew = Asys\bsys

    return BSpline(knotsNew, verticesNew)
end

expandBSpline(bs::BSpline, xin, xend) = expandBSpline(bs.knots, bs.vertices, xin, xend)

"""
    expandBSpline2dim(knotsx, knotsy, vertices, xin, xend, yin, yend)
Given a 2D BSpline by `knotsx` , `knotsy` and `vetices`
compute the expanded 2D BSpline to the limits `xin, `xend`, `yin`, `yend`.

#Examples
```julia julia-repl
julia> knotsx = knotsy = 0:0.5:10;
julia> vertices = rand(length(knotsx)*length(knotsy ));
julia> expandBSpline2dim(knotsx, knotsy, vertices, -1., 11., -1., 11.)
BSpline2dim{Float64}([-1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0  …  7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0, 10.5, 11.0, 11.5],
[-1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0  …  7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0, 10.5, 11.0, 11.5],
[-1.0029300060952653, -0.4790706367303448, 0.03153695697931029, 0.48370119176194715, 0.8342290770083228, 1.1030443114892081, 1.2623356760827873, 1.2585422976688359, 1.131589871095645, 0.9860138921208538
…  0.8984551404840291, 0.7049426036058777, 0.4552726974231055, 0.2849091663041328, 0.22446630431980158, 0.22686278460500844, 0.3058408589841374, 0.4651356516733504, 0.6732304620674567, 0.8966217568026448])
```
"""
function expandBSpline2dim(knotsx, knotsy, vertices,
                           xin, xend, yin, yend)
    nverticesx = length(knotsx)
    nverticesy = length(knotsy)

    nKnotsLeftx, nKnotsRightx = nVerticesNeededLeftRight(knotsx, xin, xend)
    knotsNewx = extendedKnots(knotsx, nKnotsLeftx, nKnotsRightx)

    nverticesxNew = length(knotsNewx)

    nKnotsLefty, nKnotsRighty = nVerticesNeededLeftRight(knotsy, yin, yend)
    knotsNewy = extendedKnots(knotsy, nKnotsLefty, nKnotsRighty)

    nverticesyNew = length(knotsNewy)

    indicesVector = reshape(1:nverticesxNew*nverticesyNew, nverticesxNew, nverticesyNew)

    yindVerticesNew = zeros(nverticesxNew*nverticesyNew)
    NNew = zeros(nverticesxNew*nverticesyNew, nverticesxNew*nverticesyNew)

    # Identify a priori known vertices

    verticesMatrix = reshape(vertices, nverticesx, nverticesy)

    w = 1e10

    for i = (nKnotsLeftx+1):(nverticesx+nKnotsLeftx)
        for j = (nKnotsLefty+1):(nverticesy+nKnotsLefty)
            locali = i - nKnotsLeftx  ;localj = j - nKnotsLefty
            yindVerticesNew[indicesVector[i,j]] = verticesMatrix[locali,localj]
            NNew[indicesVector[i,j],indicesVector[i,j]] = 1
        end
    end

    Wmat = Diagonal(fill(w, size(NNew,1)))

    Px, Py = PD2matrices2dim(nverticesxNew, nverticesyNew)
    λx = 1.
    λy = 1.

    Asys = NNew'*Wmat*NNew + λx*Px + λy*Py
    bsys = NNew'*Wmat*yindVerticesNew

    verticesNew =  Asys\bsys

    return BSpline2dim(knotsNewx, knotsNewy, verticesNew)
end

expandBSpline2dim(bs2dim::BSpline2dim, xin, xend, yin, yend) = expandBSpline2dim(bs2dim.knotsx, bs2dim.knotsy, bs2dim.vertices, xin, xend, yin, yend)

## Document
function buildSystemMatrix(N;
                            W = Wmatrix(size(N,1)),
                           Ω1 = Ω1matrix(size(N,2)),
                           Ω2 = Ω2matrix(size(N,2)),
                           Ω3 = Ω3matrix(size(N,2)))
    nvertices = size(N,2)
    D1 = D1matrix(nvertices)
    D2 = D2matrix(nvertices)
    D3 = D3matrix(nvertices)

    Asys = N'*W*N + D1'*Ω1*D1 + D2'*Ω2*D2 + D3'*Ω3*D3

    return Asys
end
## Document
function buildSystemIndVect(N, y;
                            W = Wmatrix(size(N,1)))
    bsys = N'*W*y
    return bsys
end


# test_meshwarp()
println("testing meshwarp")
img = reshape(float(collect(1:121) .% 2), 11, 11) # 11x11 checkerboard
triangles = [1 2 3]

tform = [
    1 0 0;
    0 1 0;
    0 0 1
]
src = [
    0.0 0.0;
    0.0 10.0;
    10.0 0.0
]
dst = warp_pts(tform, src)
offset = [0, 0]
img_warped, warped_offset = meshwarp(img, Array(src'), Array(dst'), Array(triangles'), offset)
# @test_approx_eq_eps img_warped img 1e-5
@test warped_offset == [0, 0]

tform = [
    1 0 0;
    0 1 0;
    0 0 1]
src = [
    0.0 0.0;
    0.0 10.0;
    10.0 0.0]
dst = warp_pts(tform, src)
offset = [5, 10]
img_warped, warped_offset = meshwarp(img, Array(src'), Array(dst'), Array(triangles'), offset)
@test img_warped ≈ zeros(11, 11) atol=1e-10
@test warped_offset == [0, 0]

tform = [1 0 0; 0 1 0; 0 0 1]
src = [0.0 0.0; 0.0 20.0; 20.0 0.0]
dst = warp_pts(tform, src)
offset = [5, 10]
img_warped, warped_offset = meshwarp(img, Array(src'), Array(dst'), Array(triangles'), offset)
@test img_warped[6, 11] == 1.0
@test warped_offset == [0, 0]

src = [2.0 2.0; 10.0 2.0; 6.0 10.0]
dst = [4.0 2.0; 8.0 2.0; 6.0 10.0]
offset = [0, 0]
img_warped, warped_offset = meshwarp(img, Array(src'), Array(dst'), Array(triangles'), offset)
@test warped_offset == [4, 2]

# test_poly2mask()
# vertices=[1.5 2.5; 2.5 1.5; 3.5 4.5]
# img, off = poly2mask(vertices[:,1],vertices[:,2])
# """ off should be [1,1] and img should be:
# false  false  false  false  false
# false   true   true  false  false
# false  false   true   true  false
# false  false  false  false  false
# """
# vertices=[1.5 2.4; 2.5 1.5; 3.5 4.5]
# img, off = poly2mask(vertices[:,1],vertices[:,2])
# """ off should be [1,1] and img should be:
# false  false  false  false  false
# false   true  false  false  false
# false  false   true  false  false
# false  false  false  false  false,
# """
# vertices=[1 2; 2.5 1.5; 3.5 4.5]
# img, off = poly2mask(vertices[:,1],vertices[:,2])
# """ corner case: off should be [1,1] and img should be:
# false   true  false  false  false
# false   true   true  false  false
# false  false   true   true  false
# false  false  false  false  false
# """
# p=plot(x=vertices[:,1],y=vertices[:,2])
# display(p)
# img,off


# Example of robot kinematics constraints reported by John Gustafson
# http://www.johngustafson.net/presentations/Multicore2016-JLG.pdf

using ConstraintPropagation

d = Domain()

@add_constraint d s2*c5*s6 - s3*c5*s6 - s4*c5*s6 + c2*c6 + c3*c6 + c4*c6 == 0.4077
@add_constraint d c1*c2*s5 + c1*c3*s5 + c1*c4*s5 + s1*c5 == 1.9115
@add_constraint d s2*s5 + s3*s5 + s4*s5 == 1.9791
@add_constraint d c1*c2 + c1*c3 + c1*c4 + c1*c2 + c1*c3 + c1*c2 == 4.0616
@add_constraint d s1*c2 + s1*c3 + s1*c4 + s1*c2 + s1*c3 + s1*c3 == 1.7172
@add_constraint d s2 + s3 + s4 + s2 + s3 + s2 == 3.9701
@add_constraint d s1^2 + c1^2 == 1
@add_constraint d s2^2 + c2^2 == 1
@add_constraint d s3^2 + c3^2 == 1
@add_constraint d s4^2 + c4^2 == 1
@add_constraint d s5^2 + c5^2 == 1
@add_constraint d s6^2 + c6^2 == 1


# Example: Use monthly average temperature/salinity fields on poseidon to compute sigma1
using Revise
using ECCOtour, MITgcmTools, MeshArrays

# other useful packages in Reemergence project
# using Statistics, PyPlot, Distributions, FFTW, LinearAlgebra, StatsBase

# get monthly average temperature and salinity
include("examples/read_TS_monavg_poseidon.jl")

p₀ = 1000 # reference pressure = 1000 dbar

# ECCOv4r4 uses approximation for pressure without any horizontal deviations.
# Can precompute pressure for each depth level.
ρ₀ = 1029 # from "data" text file in run directory
g  = 9.81 # from "data"

# depths of MITgcm levels
pathLLC = "../inputs/GRID_LLC90/" # need to input grid earlier: see example
fileZ = "RC" # "R" = radius, "C"= center of grid cell
Z = read_mdsio(pathLLC,fileZ)
Z = vec(Z)

# standard pressures via hydrostatic balance
Pa2dbar = 1/10000
p = -ρ₀ .*g .* Pa2dbar .* Z # 10000 to change Pa to dbar

# mitgcmtools.seawaterdensity is set up for a column.
# random choice of a water column.
ix = 50
iy = 400

# filter out zeroes.
θz = filter(!iszero,θ[ix,iy,:])
Sz = filter(!iszero,S[ix,iy,:])
nz = length(θz)
pz = p[1:nz]
p₀ = 1000

# ρI is in-situ density, ρR is density referenced to p0.
(ρP,ρI,ρR) = SeaWaterDensity(θz,Sz,pz,p₀)
σ₁ = ρR .- 1000

# could add loops for the entire domain.

# to do: plot it


import Random
import LinearAlgebra

function randsdp_data(seed, m, n)
    rng = Random.MersenneTwister(seed)
    # n = 15  # Instance size
    # m = 10  # Number of constraints
    # Objective function
    c_sqrt = Random.rand(rng, Float64, (n, n))
    C = c_sqrt * c_sqrt'
    # C[1, 2] *= 0.5
    # C[2, 1] = C[1, 2]
    # Generate m-dimensional feasible system
    A, b = Dict(), Dict()
    X_ = Random.randn(rng, (n, n))
    X_ = X_ * X_'
    for i in 1:m
        A[i] = Random.rand(rng, Float64, (n, n))
        A[i] = A[i] * A[i]'
        b[i] = tr(A[i] * X_)
    end
    return A, b, C
end
function randsdp_eval(A,b,C,n,m,XX)
    @show minus_rank = length([eig for eig in LinearAlgebra.eigen(XX).values if eig < -1e-10])

    @show rank = length([eig for eig in LinearAlgebra.eigen(XX).values if eig > 1e-10])

    @show tr(C * XX)
    for i in 1:m
        @show tr(A[i] * XX)-b[i]
    end

    nothing
end


# Load Julia packages (libraries) needed for clip

using StatisticalRethinking

ProjDir = @__DIR__

# Define the Stan language model

m_5_3 = "
data {
  int N;
  vector[N] D;
  vector[N] M;
  vector[N] A;
}
parameters {
  real a;
  real bA;
  real bM;
  real<lower=0> sigma;
}
model {
  vector[N] mu = a + bA * A + bM * M;
  a ~ normal( 0 , 0.2 );
  bA ~ normal( 0 , 0.5 );
  bM ~ normal( 0 , 0.5 );
  sigma ~ exponential( 1 );
  D ~ normal( mu , sigma );
}
";

# Define the SampleModel
m5_3s = SampleModel("m5.3", m_5_3);

# Input data

m5_3_data = Dict(
  "N" => size(df, 1),
  "D" => df[:, :D],
  "M" => df[:, :M],
  "A" => df[:, :A]
);

# Sample using cmdstan

rc = stan_sample(m5_3s, data=m5_3_data);

if success(rc)

  # Describe the draws

  dfa3 = read_samples(m5_3s; output_format=:dataframe)
  Particles(dfa3)

end


function MakeP_bar(A,b,C,d,XLB,XUB,y;Cuts = [])
    m,n = size(A);
    nObj = size(C)[1];
    MyModel = Model();
    @variable(MyModel, XLB[i] <= X[i=1:n] <= XUB[i])
    @variable(MyModel,Z)
    @objective(MyModel,Min,Z)
    @constraintref Ucons[1:m];
    @constraintref LAMBDAcons[1:nObj];
    for i = 1:m
        Ucons[i] = @constraint(MyModel,sum(A[i,j]*X[j] for j = 1:n) >= b[i] )
    end
    for j = 1:nObj
        LAMBDAcons[j] = @constraint(MyModel,sum(C[j,i]*X[i] for i = 1:n) + d[j] - Z <= y[j] )
    end
    for c = 1:length(Cuts)
        @constraint(MyModel,sum(Cuts[c].Coeff[i]*(sum(C[i,j]*X[j] for j = 1:n )+d[i])  for i = 1:nObj) >= Cuts[c].RHS)
    end
    return MyModel,Ucons,LAMBDAcons;

end