"""
    struct GlobalMatrices

Stores global matrices and some other things.
Everything that is very expensive to compute which stays
constant throughout the simulation should be stored here.

    GlobalMatrices(basis::Basis, filter::Filter, dimensions)

Initialize GlobalMatrices for `basis`, `filter` and `dimensions`.
"""
 struct GlobalMatrices
     normalsigns::Dict{Face, Int64}
     normalidxs::Dict{Face, Int64}
     oppositefaces::Dict{Face, Face}
     project_dofs_to_face::Dict{Face,Array{Float64,2}}
     project_flux_to_face::Dict{Face,Array{Float64,2}}
     project_dofs_from_face::Dict{Face,Array{Float64,2}}
     basissize_1d::Int64
     quadweights_nd::Array{Float64, 1}
     reference_massmatrix_face::Array{Float64,2}
     reference_massmatrix_cell::Array{Float64,2}
     reference_derivative_matrix::Array{Float64,2}

     filter_matrix::Array{Float64,2}

     function GlobalMatrices(basis::Basis, filter::Filter, dimensions)
        faces = instances(TerraDG.Face)

        normalsigns = Dict(left => -1, right => 1, top => 1, bottom => -1)
        normalidxs = Dict(left => 1, right => 1, top => 2, bottom => 2)
        oppositefaces = Dict(left => right, right => left, top => bottom, bottom => top)

        project_dofs_to_face = Dict(f => face_projection_matrix(basis, f) for f in faces)
        project_flux_to_face = Dict(f => face_projection_matrix(basis, f) for f in faces)
        project_dofs_from_face = Dict(f => transpose(face_projection_matrix(basis, f)) for f in faces)

        basissize_1d = size(basis,1)
        basissize_nd = length(basis)
        quadweights_nd = basis.quadweights;

        reference_massmatrix_cell = massmatrix(basis, dimensions)
        reference_massmatrix_face = massmatrix(basis, dimensions - 1)
        reference_derivative_matrix = derivativematrix(basis)

        filter_matrix = evaluate_filter_matrix(filter, basis)

         new(normalsigns, normalidxs, oppositefaces, project_dofs_to_face,
         project_flux_to_face, project_dofs_from_face, basissize_1d, quadweights_nd,
         reference_massmatrix_face,
         reference_massmatrix_cell, reference_derivative_matrix,
         filter_matrix)
     end
 end


# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

@userplot VarPlane

@recipe function f(vp::VarPlane;
                   theta=0, phi=90,
                   ptol=1e-6, dtol=0.5,
                   nangs=50, nlags=20, maxlag=nothing,
                   showrange=true, varmodel=GaussianVariogram)
  # sanity checks
  @assert 0 ≤ theta ≤ 360 "theta must lie in [0,360]"
  @assert -90 ≤ phi ≤ 90 "phi must lie in [-90,90]"

  # get inputs
  spatialdata = vp.args[1]
  var₁ = vp.args[2]
  var₂ = length(vp.args) == 3 ? vp.args[3] : var₁
  N = ndims(spatialdata)

  # basis for variogram plane
  if N == 2
    planes = [spatialdata]
    u, v = (1.,0.), (0.,1.)
  elseif N == 3
    θ = deg2rad(theta); ϕ = deg2rad(phi)
    n = (cos(ϕ)cos(θ), cos(ϕ)sin(θ), sin(ϕ))

    planes = partition(spatialdata, PlanePartitioner(n, tol=ptol))
    u, v = plane_basis(n)
  else
    @error "variogram plane only supported in 2D or 3D"
  end

  # polar coordinates and variogram values
  rs = Vector{Float64}(undef, nlags)
  θs = range(0, stop=π, length=nangs)
  zs = Matrix{Float64}(undef, nlags, nangs)

  # ranges for each angle
  ls = Vector{Float64}(undef, nangs)

  # loop over half of the plane
  for (j, θ) in Iterators.enumerate(θs)
    dir = ntuple(i -> cos(θ)*u[i] + sin(θ)*v[i], N)

    dpart = DirectionPartitioner(dir, tol=dtol)

    # compute directional variogram across planes
    plane, _ = iterate(planes)
    γ = EmpiricalVariogram(partition(plane, dpart), var₁, var₂;
                           nlags=nlags, maxlag=maxlag)
    for plane in Iterators.drop(planes, 1)
      γplane = EmpiricalVariogram(partition(plane, dpart), var₁, var₂;
                                  nlags=nlags, maxlag=maxlag)
      merge!(γ, γplane)
    end

    rs, ys, _ = values(γ)

    # clean NaN values (i.e. empty bins)
    isnan(ys[1]) && (ys[1] = 0)
    for i in 2:length(ys)
      isnan(ys[i]) && (ys[i] = ys[i-1])
    end

    # save variogram for this angle
    zs[:,j] = ys

    if showrange
      γtheo = fit(varmodel, γ)
      ls[j] = range(γtheo)
    end
  end

  # exploit symmetry
  θs = range(0, stop=2π, length=2nangs)
  zs = hcat(zs, zs)

  # show variables in plot title
  ptitle = var₁ == var₂ ? "$var₁" : "$var₁-$var₂"

  projection --> :polar
  ylim --> (0, Inf)
  title --> ptitle

  # plot variogram plane
  @series begin
    seriestype --> :heatmap
    color --> :curl

    θs, rs, zs
  end

  # plot ranges
  if showrange
    @series begin
      seriestype --> :path
      linecolor --> :black
      label --> "range"

      θs, ls
    end
  end
end

function plane_basis(normal::NTuple{3,T}) where {T<:Real}
  # normalize input
  n = normal ./ sqrt(sum(normal[i]^2 for i in 1:3))

  # find last non-zero component
  idx = -1
  for (i, c) in enumerate(reverse(n))
    if c != 0
      idx = length(n) - i + 1
      break
    end
  end

  @assert idx > 0 "invalid normal vector"

  # first basis vector (perturb and subtract projection)
  u = ntuple(i -> i == idx%3 + 1 ? n[i] + one(T) : n[i], 3)
  l = sum(u[i]*n[i] for i in 1:3)
  u = ntuple(i -> u[i] - l*n[i], 3)

  # second basis vector (cross product)
  nx, ny, nz = n
  ux, uy, uz = u
  v = (ny*uz - nz*uz, nz*ux - nx*uz, nx*uy - ny*ux)

  # normalize output
  u = u ./ sqrt(sum(u[i]^2 for i in 1:3))
  v = v ./ sqrt(sum(v[i]^2 for i in 1:3))

  u, v
end


using Test;

function run(fname, version=1)
    f = open(fname, "r")
    lines = readlines(f)
    close(f)

    result = 0
    for l in lines
        if version == 1
            result += compute(l)
        elseif version == 2
            result += compute2(l)
        end
    end

    println("[Part ", version, "] Result = ", result)
    return result
end

function add(a, b)
    return a + b
end

function multiply(a, b)
    return a * b
end

funcs = Dict()
funcs["+"] = add
funcs["*"] = multiply

# WARNING: This is special math
function compute(s, version=1)
    if occursin('(', s)
        # if you find a pair of parenthesis, compute what's inside recursively
        start_idx = findlast(isequal('('), s)
        end_idx = findnext(isequal(')'), s, start_idx)

        pre = s[1:start_idx - 1]
        mid = compute(s[start_idx + 1:end_idx - 1], version)
        post = s[end_idx + 1:length(s)]
        return compute(pre * string(mid) * post, version)
    else
        parts = split(s)
        while length(parts) > 1
            start = 1
            if version == 2 && in("+", parts)
                # "Advanced Math": plus is resolved before multiply
                start = findfirst(isequal("+"), parts) - 1
            end
            val = resolve(parts[start:start + 2])
            parts = [parts[1:start - 1]..., string(val), parts[start + 3:length(parts)]...]
        end

        total = parse(Int, parts[1])
        return total
    end
end

function resolve(list)
    # println(list)
    # given a list of 3 items (value1, operation, value2), reduces to 1 item (output value)
    v1 = parse(Int, list[1])
    v2 = parse(Int, list[3])
    op = add
    if list[2] == "*"
        op = multiply
    end

    return op(v1, v2)
end

function compute2(s)
    return compute(s, 2)
end

# @test compute("1 + 2 * 3 + 4 * 5 + 6") == 71
# @test compute("2 * 3 + (4 * 5)") == 26
# @test compute("5 + (8 * 3 + 9 + 3 * 4 * 3)") == 437
# @test compute("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))") == 12240
# @test compute("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2") == 13632

# run("18.txt")

@test resolve(["1", "+", "3"]) == 4
@test resolve(["1", "*", "3"]) == 3

@test compute2("1 + 2 * 3 + 4 * 5 + 6") == 231
@test compute2("1 + (2 * 3) + (4 * (5 + 6))") == 51
@test compute2("2 * 3 + (4 * 5)") == 46
@test compute2("5 + (8 * 3 + 9 + 3 * 4 * 3)") == 1445
@test compute2("5 * 9 * (7 * 3 * 3 + 9 * 3 + (8 + 6 * 4))") == 669060
@test compute2("((2 + 4 * 9) * (6 + 9 * 8 + 6) + 6) + 2 + 4 * 2") == 23340

run("18.txt", 2)


#####
##### A simple minmax player to be used as a baseline
#####

"""
A simple implementation of the minmax tree search algorithm, to be used as
a baseline against AlphaZero. Heuristic board values are provided by the
[`GameInterface.heuristic_value`](@ref) function.
"""
module MinMax

import ..GI, ..GameInterface, ..AbstractPlayer, ..think

function current_player_value(white_reward, white_playing) :: Float64
  if iszero(white_reward)
    return 0.
  else
    v = Inf * sign(white_reward)
    return white_playing ? v : - v
  end
end

# Return the value of a state for the player playing
function value(game, depth)
  wr = GI.white_reward(game)
  wp = GI.white_playing(game)
  if isnothing(wr)
    if depth == 0
      return GI.heuristic_value(game)
    else
      return maximum(qvalue(game, a, depth)
        for a in GI.available_actions(game))
    end
  else
    return current_player_value(wr, wp)
  end
end

function qvalue(game, action, depth)
  @assert isnothing(GI.white_reward(game))
  next = copy(game)
  GI.play!(next, action)
  pswitch = GI.white_playing(game) != GI.white_playing(next)
  nextv = value(next, depth - 1)
  return pswitch ? - nextv : nextv
end

minmax(game, actions, depth) = argmax([qvalue(game, a, depth) for a in actions])

"""
    MinMax.Player{Game} <: AbstractPlayer{Game}

A stochastic minmax player, to be used as a baseline.

    MinMax.Player{Game}(;depth, τ=0.)

The minmax player explores the game tree exhaustively at depth `depth`
to build an estimate of the Q-value of each available action. Then, it
chooses an action as follows:

- If there are winning moves (with value `Inf`), one of them is picked
  uniformly at random.
- If all moves are losing (with value `-Inf`), one of them is picked
  uniformly at random.

Otherwise,

- If the temperature `τ` is zero, a move is picked uniformly among those
  with maximal Q-value (there is usually only one choice).
- If the temperature `τ` is nonzero, the probability of choosing
  action ``a`` is proportional to ``e^{\\frac{q_a}{Cτ}}`` where ``q_a`` is the
  Q value of action ``a`` and ``C`` is the maximum absolute value of all
  finite Q values, making the decision invariant to rescaling of
  [`GameInterface.heuristic_value`](@ref).
"""
struct Player{G} <: AbstractPlayer{G}
  depth :: Int
  τ :: Float64
  Player{G}(;depth, τ=0.) where G = new{G}(depth, τ)
end

function think(p::Player, game, turn=nothing)
  actions = GI.available_actions(game)
  n = length(actions)
  qs = [qvalue(game, a, p.depth) for a in actions]
  winning = findall(==(Inf), qs)
  if isempty(winning)
    notlosing = findall(>(-Inf), qs)
    best = argmax(qs)
    if isempty(notlosing)
      π = ones(n)
    elseif iszero(p.τ)
      π = zeros(n)
      all_best = findall(==(qs[best]), qs)
      π[all_best] .= 1.
    else
      qmax = qs[best]
      @assert qmax > -Inf
      C = maximum(abs(qs[a]) for a in notlosing) + eps()
      π = exp.((qs .- qmax) ./ C)
      π .^= (1 / p.τ)
    end
  else
    π = zeros(n)
    π[winning] .= 1.
  end
  π ./= sum(π)
  return actions, π
end

end


# Planetary impact cross section; "critical" B
# derived from conservation of energy and angular momentum
# in hyperbolic close encounter
#  = what impact parameter B corresponds to a grazing impact,
# given planet radius and mass, and asymptotic vel v_infty
# if actual B is equal or less to this, then impact happens
# μ_P: planetary gravitational parameter (au^3/day^2)
# R_P: planetary radius (au)
# vinf: asymptotic inbound velocity (au/day)
# output is in planet radii
function crosssection(μ_P, R_P, vinf)
    return sqrt( 1 + (2μ_P)/(R_P*(vinf)^2) )
end

# computes Öpik's coordinates of impact parameter vector B in hyperbolic planetary close encounter
# following Farnocchia et al (2019)
# xae: asteroid's geocentric position/velocity vector at closest approach in au,au/day
# xes: planet's heliocentric position/velocity vector at asteroid's closest approach in au,au/day
function bopik(xae, xes)
    μ_E = PlanetaryEphemeris.μ[ea]
    # asteroid geocentric range at 2029 closest approach (au)
    rae = sqrt(xae[1]^2 + xae[2]^2 + xae[3]^2)
    # osculating semimajor axis at closest approach (negative since hyperbolic)
    a = semimajoraxis(xae..., μ_E, 0.0)
    # asymptotic inbound velocity v_\infty (au/day)
    v_infty = sqrt(μ_E/(-a))
    # h = r × v
    hvec = cross(xae[1:3], xae[4:6])
    h = sqrt(hvec[1]^2 + hvec[2]^2 + hvec[3]^2)
    # Laplace-Runge-Lenz (eccentricity) vector
    # \vec{e} = (\vec{v} × \vec{h})/μ - \vec{r}/r
    evec = cross(xae[4:6], hvec)/μ_E - xae[1:3]/rae
    # osculating eccentricity
    e = eccentricity(xae..., μ_E, 0.0)
    # osculating inclination
    i = inclination(xae...)
    # periapsis position (unit vector)
    P_v = evec./e
    # periapsis velocity (unit vector)
    Q_v = cross(hvec, P_v)./h
    #inbound asymptote direction
    S_v = (P_v + (sqrt(e^2 - 1))Q_v)/e
    # B-vector: "vector from the planet center to the intersection between the B-plane and the asymptote"
    Bvec = cross(S_v, hvec)./v_infty
    # Earth impact cross section ("critical B")
    b_E = crosssection(μ_E, RE/au, v_infty)
    # @show b_E
    # Earth's heliocentric velocity at CA
    v_pl = xes[4:6]
    #ξ-axis is essentially the MOID (Valsecchi et al, 2003)
    ξ_v_unnormalized = cross(v_pl, S_v)
    ξ_v_norm = sqrt(ξ_v_unnormalized[1]^2 + ξ_v_unnormalized[2]^2 + ξ_v_unnormalized[3]^2)
    ξ_v = ξ_v_unnormalized./ξ_v_norm
    # @show ξ_v
    #ζ-axis: delay/advance in CA time (Valsecchi et al, 2003)
    ζ_v = -cross(S_v, ξ_v)
    B_dot_ξ = dot(Bvec, ξ_v)
    B_dot_ζ = dot(Bvec, ζ_v)

    # computation of U vector
    # res = xes[1:3] # Earth's heliocentric position, au
    # res_norm = norm(res) # Earth's heliocentric range, au
    # res_unit = res/res_norm # Earth's heliocentric radius unit vector
    # @show res_norm
    ves = v_pl*au/daysec # Earth's velocity, km/s
    ves_norm = norm(ves) # Earth's speed, km/s
    ves_unit = ves/ves_norm # Earth's velocity unit vector
    # X-Y angle (degrees)
    # @show rad2deg(acos(dot(ves_unit, res_unit)))
    # angle between Y-axis and \vec{U}
    cosθ = dot(S_v, ves_unit)
    # @show cosθ()
    v_infty_kms = v_infty*au/daysec # asteroid unperturbed speed, km/sec
    # @show v_infty_kms()
    # The norm of \vec{U} in appropriate units
    U_unit = (2pi*au)/(yr*daysec) # 1 U in km/s
    U_norm = v_infty_kms/U_unit
    # U_y
    U_y = U_norm*cosθ
    # @show U_y U_norm

    return (ξ=B_dot_ξ*au/RE, ζ=B_dot_ζ*au/RE, U=(y=U_y, norm=U_norm), b=b_E)
end

# computes Valsecchi circle associated to a mean motion resonance, following Valsecchi et al (2003)
# Returns radius R (au) and zeta-axis coordinate D (au)
#a: asteroid heliocentric semimajor axis (au)
#e: asteroid heliocentric eccentricity
#i: asteroid heliocentric inclination, ecliptic (rad)
#k/h; h heliocentric revolutions of asteroid per k heliocentric revolutions of Earth
# m_pl: planet mass normalized to Sun's mass, equal to Earth mass in solar masses by default
function valsecchi_circle(a, e, i, k, h; m_pl=3.003489614915764e-6)
    # U_x = sqrt( 2 - (1/a) - a*(1-(e^2)) ) # TODO: CHECK SIGN
    U_y = sqrt( a*(1-(e^2)) )*cos(i) - 1
    # U_z = sqrt( a*(1-(e^2)) )*sin(i) # TODO: CHECK SIGN
    U_norm = sqrt( 3 - (1/a) - 2*sqrt(a*(1-(e^2)))*cos(i) )
    # expression below should be equal to asteroid heliocentric elliptic semimamajor axis in au units
    #@show 1/(1-U_^2-2U_y)
    return valsecchi_circle(U_y, U_norm, k, h, m_pl=m_pl)
end

# computes Valsecchi circle associated to a mean motion resonance, following Valsecchi et al (2003)
# Returns radius R (au) and zeta-axis coordinate D (au)
# U_y: Y-component of unperturbed planetocentric velocity (Y-axis coincides with the direction of motion of the planet)
# U_norm: Euclidean norm of unperturbed planetocentric velocity
# both U_y, U_norm are in units such that the heliocentric velocity of the planet is 1
#k/h; h heliocentric revolutions of asteroid per k heliocentric revolutions of Earth
# m_pl: planet mass normalized to Sun's mass, equal to Earth mass in solar masses by default
# a_pl: planetary heliocentric semimajor axis in au; default value is 1
function valsecchi_circle(U_y, U_norm, k, h; m_pl=3.003489614915764e-6, a_pl=1.0)
    a0p = a_pl*(k/h)^(2/3)
    cosθ = U_y/U_norm
    # @show cosθ, U_norm, U_y
    sinθ = sin(acos(cosθ)) #sqrt(1-cosθ^2) #  # TODO: CHECK SIGN
    cosθ0p = (1-(U_norm^2)-(1/a0p))/(2U_norm)
    sinθ0p = sin(acos(cosθ0p))
    # c = m/U^2 (Valsecchi et al, 2003, Sec. 2.3, first sentence below Eq. 3)
    c = m_pl/(U_norm^2)
    R0 = abs( c*sinθ0p/(cosθ0p-cosθ) )
    D0 = c*sinθ/(cosθ0p-cosθ)
    return R0, D0
end


using SemiDiscretizationMethod
using Test

function createHayesProblem(a,b)
    AMx =  ProportionalMX(a*ones(1,1));
    τ1=1.
    BMx1 = DelayMX(τ1,b*ones(1,1));
    cVec = Additive(ones(1))
    LDDEProblem(AMx,[BMx1],cVec)
end

function createMathieuProblem(δ,ε,b0,a1;T=2π)
    AMx =  ProportionalMX(t->@SMatrix [0. 1.; -δ-ε*cos(2π/T*t) -a1]);
    τ1=2π # if function is needed, the use τ1 = t->foo(t)
    BMx1 = DelayMX(τ1,t->@SMatrix [0. 0.; b0 0.]);
    cVec = Additive(t->@SVector [0.,sin(4π/T*t)])
    LDDEProblem(AMx,[BMx1],cVec)
end

function tests()
    @testset "Testing the SemiDiscretizationMethod package with the examples" begin
		#Hayes
        @test begin
            hayes_lddep=createHayesProblem(-1.,-1.); # LDDE problem for Hayes equation
			method=SemiDiscretization(1,0.1) # 3rd order semi discretization with Δt=0.1
			τmax=1. # the largest τ of the system
			mapping=DiscreteMapping(hayes_lddep,method,τmax,n_steps=1,calculate_additive=true); #The discrete mapping of the system
			spectralRadiusOfMapping(mapping)
			fixPointOfMapping(mapping)
            true
        end
		#Delay Mathieu
        @test begin
            τmax=2π # the largest τ of the system
			T=2π #Principle period of the system (sin(t)=sin(t+P))
			mathieu_lddep=createMathieuProblem(3.,2.,-0.15,0.1,T=T); # LDDE problem for Hayes equation
			method=SemiDiscretization(1,0.01) # 3rd order semi discretization with Δt=0.1
			# if P = τmax, then n_steps is automatically calculated
			mapping=DiscreteMapping(mathieu_lddep,method,τmax,
		    n_steps=Int((T+100eps(T))÷method.Δt),calculate_additive=true); #The discrete mapping of the system

			spectralRadiusOfMapping(mapping); # spectral radius ρ of the mapping matrix (ρ>1 unstable, ρ<1 stable)
			# spectralRadiusOfMapping(mapping) = 0.5131596340374617
			fp=fixPointOfMapping(mapping); # stationary solution of the hayes equation (equilibrium position)
			true
        end
    end
end
tests()


using WoodburyMatrices

d  = 2+rand(5)
dl = -rand(4)
du = -rand(4)
M = Tridiagonal(dl, d, du)
F = lufact(M)
U = sprand(5,2,0.2)
V = sprand(2,5,0.2)
C = rand(2,2)
W = Woodbury(F, U, C, V)

src = rand(5, 8)
@test_approx_eq W\src AxisAlgorithms.A_ldiv_B_md(W, src, 1)
src = rand(5, 5, 5)
for dim = 1:3
    dest1 = mapslices(x->W\x, copy(src), dim)
    dest2 = similar(src)
    AxisAlgorithms.A_ldiv_B_md!(dest2, W, src, dim)
    @test_approx_eq dest1 dest2
end


using StatisticalNonlocality, LinearAlgebra, FFTW, SparseArrays, JLD2
import StatisticalNonlocality: chebyshev, fourier_nodes, fourier_wavenumbers
import StatisticalNonlocality: droprelativezeros!

# transition rate
γ = 1e0  # 1e0 # 1e2
# diffusivity
κ = 1e-2 # 1e0 # 1e-2

# minimal configuration is M = 64, N = 8
N = 8 * 4
M = 8 * 4

# The first step is to build the block operators
a, b = 0, 4π
k = fourier_wavenumbers(N, L = b - a)
x = fourier_nodes(N, a = a, b = b)

ℱ = fft(I + zeros(N, N), 1)
ℱ⁻¹ = ifft(I + zeros(N, N), 1)
Dx = real.(ℱ⁻¹ * Diagonal(im .* k) * ℱ)
Dy = copy(Dx)
x = reshape(x, (N, 1))
y = reshape(x, (1, N))
∂y = kron(Dy, I + zeros(N, N))
∂x = kron(I + zeros(N, N), Dx)
Δ = ∂x^2 + ∂y^2

# convenience functions
advection_operator(ψ, ∂x, ∂y) = Diagonal(∂y * ψ[:]) * ∂x + Diagonal(-∂x * ψ[:]) * ∂y

makelocal(E) = Array(Diagonal(sum(E, dims = 2)[:]))

function localdiffusivity(E, N, M)
    local localdiff = makelocal(E)
    localdiff = [localdiff[i, i] for i = 1:size(E)[1]]
    local localdiff = reshape(localdiff, (N, M ))
    return localdiff
end

function avglocaldiffusivity(E, N, M)
    local localdiff = localdiffusivity(E, N, M)
    local localdiff = sum(localdiff, dims = 1)[:] ./ N # zero'th mode
    return localdiff
end

# the local diffusivity is proportional to 1/γ
function grabdiagonal(A)
    MM = minimum(size(A))
    diagA = zeros(MM)
    for i in 1:MM
        diagA[i] = A[i, i]
    end
    return diagA
end

# random phase wavenumbers
wavemax = 0.5 # change to 3 when automated
𝓀 = collect(-wavemax:0.5:wavemax)
𝓀ˣs = reshape(𝓀, (length(𝓀), 1))
𝓀ʸs = reshape(𝓀, (1, length(𝓀)))
A = @. 0.1 * (𝓀ˣs * 𝓀ˣs + 𝓀ʸs * 𝓀ʸs)^(-11 / 12)
A[A.==Inf] .= 0.0

tic = Base.time()
for (i, 𝓀ˣ) in enumerate(𝓀ˣs), (j, 𝓀ʸ) in enumerate(𝓀ʸs)
    filename = "lots_of_diffusivities_kx_" * string(𝓀ˣ) * "ky_" * string(𝓀ʸ) * ".jld2"
    println("-----------")
    println("currently on wavenumber 𝓀ˣ=", 𝓀ˣ, ", 𝓀ʸ=", 𝓀ʸ)
    println("The amplitude is ", A[i, j])
    println("-----------")
    U₀ = A[i, j]

    ψ¹ = @. U₀ * sin(𝓀ˣ * x + 𝓀ʸ * y)
    ψ² = @. U₀ * cos(𝓀ˣ * x + 𝓀ʸ * y)
    ψ³ = @. -U₀ * sin(𝓀ˣ * x + 𝓀ʸ * y)
    ψ⁴ = @. -U₀ * cos(𝓀ˣ * x + 𝓀ʸ * y)

    A¹ = advection_operator(ψ¹, ∂x, ∂y)
    A² = advection_operator(ψ², ∂x, ∂y)
    A³ = advection_operator(ψ³, ∂x, ∂y)
    A⁴ = advection_operator(ψ⁴, ∂x, ∂y)

    L = [(γ*I-κ*Δ) 0.0*I 0.5*A¹; 0.0*I (γ*I-κ*Δ) -0.5*A²; A¹ -A² (2*γ*I-κ*Δ)]
    G = L \ I

    # Create the block linear operator for components of the diffusivity tensor
    u¹ = Diagonal(∂y * ψ¹[:])
    u² = Diagonal(∂y * ψ²[:])
    v¹ = Diagonal(-∂x * ψ¹[:])
    v² = Diagonal(-∂x * ψ²[:])

    U = -0.5 * [u¹; u²; 0 * I]
    Uᵀ = [u¹ u² 0 * I]
    V = -0.5 * [v¹; v²; 0 * I]
    Vᵀ = [v¹ v² 0 * I]

    EF¹¹ = Uᵀ * G * U # flux in x due to gradients in x
    EF¹² = Uᵀ * G * V # flux in x due to gradients in z
    EF²¹ = Vᵀ * G * U # flux in z due to gradients in x
    EF²² = Vᵀ * G * V # flux in z due to gradients in z

    E = [EF¹¹ EF¹²; EF²¹ EF²²]

    λE = eigvals(E)
    bools = real.(λE) .≤ eps(1e3 * maximum(abs.(λE)))
    println(sum(bools) == length(λE))

    # OUTPUT
    # always overwrite
    data_directory = "data"
    mkpath(data_directory)
    filepath = data_directory * "/" * filename
    if isfile(filepath)
        rm(filepath)
    end
    ##
    file = jldopen(filepath, "a+")
    # diffusivity
    groupname = "diffusivity"
    JLD2.Group(file, groupname)
    file[groupname]["K11"] = EF¹¹
    file[groupname]["K12"] = EF¹²
    file[groupname]["K21"] = EF²¹
    file[groupname]["K22"] = EF²²
    # parameters
    groupname = "parameters"
    JLD2.Group(file, groupname)
    file[groupname]["γ"] = γ
    file[groupname]["κ"] = κ
    # grid
    groupname = "grid"
    JLD2.Group(file, groupname)
    file[groupname]["x"] = x
    file[groupname]["y"] = y
    # transition matrix
    groupname = "transition"
    JLD2.Group(file, groupname)

    T = [-1 1 0 0; 0 -1 1 0; 0 0 -1 1; 1 0 0 -1]
    sT = (T + T') / 2

    file[groupname]["T"] = γ * sT
    # streamfunction
    groupname = "streamfunction"
    JLD2.Group(file, groupname)
    file[groupname]["ψ¹"] = ψ¹
    file[groupname]["ψ²"] = ψ²
    file[groupname]["ψ³"] = ψ³
    file[groupname]["ψ⁴"] = ψ⁴
    # velocities
    groupname = "velocities"
    JLD2.Group(file, groupname)
    file[groupname]["u¹"] = reshape(grabdiagonal(u¹), size(ψ¹))
    file[groupname]["u²"] = reshape(grabdiagonal(u²), size(ψ¹))
    file[groupname]["v¹"] = reshape(grabdiagonal(v¹), size(ψ¹))
    file[groupname]["v²"] = reshape(grabdiagonal(v²), size(ψ¹))

    # local diffusivities
    groupname = "localdiffusivity"
    JLD2.Group(file, groupname)
    file[groupname]["κ11"] = localdiffusivity(EF¹¹, N, N)
    file[groupname]["κ12"] = localdiffusivity(EF¹², N, N)
    file[groupname]["κ21"] = localdiffusivity(EF²¹, N, N)
    file[groupname]["κ22"] = localdiffusivity(EF²², N, N)

    # zonally averaged local diffusivities
    groupname = "averagelocaldiffusivity"
    JLD2.Group(file, groupname)
    file[groupname]["κ11"] = avglocaldiffusivity(EF¹¹, N, N)
    file[groupname]["κ12"] = avglocaldiffusivity(EF¹², N, N)
    file[groupname]["κ21"] = avglocaldiffusivity(EF²¹, N, N)
    file[groupname]["κ22"] = avglocaldiffusivity(EF²², N, N)

    close(file)

end
toc = Base.time()

println("The time for the simulation is ", (toc - tic)/60, " minutes ")


# lex indicates if its lexigraphical (i.e., x, y) or reverse (y, x)
# If in lexigraphical order the coefficient vector's entries
# corrrespond to the following basis polynomials:
# [T0(x) * T0(y), T1(x) * T0(y), T0(x) * T1(y), T2(x) * T0(y), T1(x) * T1(y), T0(x) * T2(y), ...]
# else, if not in lexigraphical order:
# [T0(x) * T0(y), T0(x) * T1(y), T1(x) * T0(y), T0(x) * T2(y), T1(x) * T1(y), T2(x) * T0(y), ...]
"""
Pre-plan an Inverse Padua Transform.
"""
struct IPaduaTransformPlan{lex,IDCTPLAN,T}
    cfsmat::Matrix{T}
    idctplan::IDCTPLAN
end

IPaduaTransformPlan(cfsmat::Matrix{T},idctplan,::Type{Val{lex}}) where {T,lex} =
    IPaduaTransformPlan{lex,typeof(idctplan),T}(cfsmat,idctplan)

"""
Pre-plan an Inverse Padua Transform.
"""
function plan_ipaduatransform!(::Type{T},N::Integer,lex) where T
    n=Int(cld(-3+sqrt(1+8N),2))
    if N ≠ div((n+1)*(n+2),2)
        error("Padua transforms can only be applied to vectors of length (n+1)*(n+2)/2.")
    end
    IPaduaTransformPlan(Array{T}(undef,n+2,n+1),FFTW.plan_r2r!(Array{T}(undef,n+2,n+1),FFTW.REDFT00),lex)
end

plan_ipaduatransform!(::Type{T},N::Integer) where {T} = plan_ipaduatransform!(T,N,Val{true})
plan_ipaduatransform!(v::AbstractVector{T},lex...) where {T} = plan_ipaduatransform!(eltype(v),length(v),lex...)

function *(P::IPaduaTransformPlan,v::AbstractVector{T}) where T
    cfsmat=trianglecfsmat(P,v)
    n,m=size(cfsmat)
    rmul!(view(cfsmat,:,2:m-1),0.5)
    rmul!(view(cfsmat,2:n-1,:),0.5)
    tensorvals=P.idctplan*cfsmat
    paduavec!(v,P,tensorvals)
end

ipaduatransform!(v::AbstractVector,lex...) = plan_ipaduatransform!(v,lex...)*v
"""
Inverse Padua Transform maps the 2D Chebyshev coefficients to the values of the interpolation polynomial at the Padua points.
"""
ipaduatransform(v::AbstractVector,lex...) = plan_ipaduatransform!(v,lex...)*copy(v)

"""
Creates ``(n+2)x(n+1)`` Chebyshev coefficient matrix from triangle coefficients.
"""
function trianglecfsmat(P::IPaduaTransformPlan{true},cfs::AbstractVector)
    N=length(cfs)
    n=Int(cld(-3+sqrt(1+8N),2))
    cfsmat=fill!(P.cfsmat,0)
    m=1
    for d=1:n+1
        @inbounds for k=1:d
            j=d-k+1
            cfsmat[k,j]=cfs[m]
            if m==N
                return cfsmat
            else
                m+=1
            end
        end
    end
    return cfsmat
end

function trianglecfsmat(P::IPaduaTransformPlan{false},cfs::AbstractVector)
    N=length(cfs)
    n=Int(cld(-3+sqrt(1+8N),2))
    cfsmat=fill!(P.cfsmat,0)
    m=1
    for d=1:n+1
        @inbounds for k=d:-1:1
            j=d-k+1
            cfsmat[k,j]=cfs[m]
            if m==N
                return cfsmat
            else
                m+=1
            end
        end
    end
    return cfsmat
end

"""
Vectorizes the function values at the Padua points.
"""
function paduavec!(v,P::IPaduaTransformPlan,padmat::Matrix)
    n=size(padmat,2)-1
    N=(n+1)*(n+2)
    if iseven(n)>0
        d=div(n+2,2)
        m=0
        @inbounds for i=1:n+1
            v[m+1:m+d]=view(padmat,1+mod(i,2):2:n+1+mod(i,2),i)
            m+=d
        end
    else
        @inbounds v[:]=view(padmat,1:2:N-1)
    end
    return v
end

"""
Pre-plan a Padua Transform.
"""
struct PaduaTransformPlan{lex,DCTPLAN,T}
    vals::Matrix{T}
    dctplan::DCTPLAN
end

PaduaTransformPlan(vals::Matrix{T},dctplan,::Type{Val{lex}}) where {T,lex} =
    PaduaTransformPlan{lex,typeof(dctplan),T}(vals,dctplan)

"""
Pre-plan a Padua Transform.
"""
function plan_paduatransform!(::Type{T},N::Integer,lex) where T
    n=Int(cld(-3+sqrt(1+8N),2))
    if N ≠ ((n+1)*(n+2))÷2
        error("Padua transforms can only be applied to vectors of length (n+1)*(n+2)/2.")
    end
    PaduaTransformPlan(Array{T}(undef,n+2,n+1),FFTW.plan_r2r!(Array{T}(undef,n+2,n+1),FFTW.REDFT00),lex)
end

plan_paduatransform!(::Type{T},N::Integer) where {T} = plan_paduatransform!(T,N,Val{true})
plan_paduatransform!(v::AbstractVector{T},lex...) where {T} = plan_paduatransform!(eltype(v),length(v),lex...)

function *(P::PaduaTransformPlan,v::AbstractVector{T}) where T
    N=length(v)
    n=Int(cld(-3+sqrt(1+8N),2))
    vals=paduavalsmat(P,v)
    tensorcfs=P.dctplan*vals
    m,l=size(tensorcfs)
    rmul!(tensorcfs,T(2)/(n*(n+1)))
    rmul!(view(tensorcfs,1,:),0.5)
    rmul!(view(tensorcfs,:,1),0.5)
    rmul!(view(tensorcfs,m,:),0.5)
    rmul!(view(tensorcfs,:,l),0.5)
    trianglecfsvec!(v,P,tensorcfs)
end

paduatransform!(v::AbstractVector,lex...) = plan_paduatransform!(v,lex...)*v
"""
Padua Transform maps from interpolant values at the Padua points to the 2D Chebyshev coefficients.
"""
paduatransform(v::AbstractVector,lex...) = plan_paduatransform!(v,lex...)*copy(v)

"""
Creates ``(n+2)x(n+1)`` matrix of interpolant values on the tensor grid at the ``(n+1)(n+2)/2`` Padua points.
"""
function paduavalsmat(P::PaduaTransformPlan,v::AbstractVector)
    N=length(v)
    n=Int(cld(-3+sqrt(1+8N),2))
    vals=fill!(P.vals,0.)
    if iseven(n)>0
        d=div(n+2,2)
        m=0
        @inbounds for i=1:n+1
            vals[1+mod(i,2):2:n+1+mod(i,2),i]=view(v,m+1:m+d)
            m+=d
        end
    else
        @inbounds vals[1:2:end]=view(v,:)
    end
    return vals
end

"""
Creates length ``(n+1)(n+2)/2`` vector from matrix of triangle Chebyshev coefficients.
"""
function trianglecfsvec!(v,P::PaduaTransformPlan{true},cfs::Matrix)
    m=size(cfs,2)
    l=1
    for d=1:m
        @inbounds for k=1:d
            j=d-k+1
            v[l]=cfs[k,j]
            l+=1
        end
    end
    return v
end

function trianglecfsvec!(v,P::PaduaTransformPlan{false},cfs::Matrix)
    m=size(cfs,2)
    l=1
    for d=1:m
        @inbounds for k=d:-1:1
            j=d-k+1
            v[l]=cfs[k,j]
            l+=1
        end
    end
    return v
end

"""
Returns coordinates of the ``(n+1)(n+2)/2`` Padua points.
"""
function paduapoints(::Type{T}, n::Integer) where T
    N=div((n+1)*(n+2),2)
    MM=Matrix{T}(undef,N,2)
    m=0
    delta=0
    NN=fld(n+2,2)
    @inbounds for k=n:-1:0
        if isodd(n)>0
            delta=mod(k,2)
        end
        @inbounds for j=NN+delta:-1:1
            m+=1
            MM[m,1]=sinpi(T(k)/n-T(0.5))
            if isodd(n-k)>0
                MM[m,2]=sinpi((2j-one(T))/(n+1)-T(0.5))
            else
                MM[m,2]=sinpi(T(2j-2)/(n+1)-T(0.5))
            end
        end
    end
    return MM
end

paduapoints(n::Integer) = paduapoints(Float64,n)


"""
# Provide implementation of default layers

Provided layers
- DenseLayer
- DenseNoBiasLayer
- VectorFunctionLayer
"""

#using Random, Zygote
#import ..Utils
#import Base.size

# ------------------------------------------------------------------------------
# DenseLayer layer

"""
   DenseLayer

Representation of a layer in the network

# Fields:
* `w`:  Weigths matrix with respect to the input from previous layer or data (n x n pr. layer)
* `wb`: Biases (n)
* `f`:  Activation function
* `df`: Derivative of the activation function
"""
mutable struct DenseLayer <: AbstractLayer
     w::Array{Float64,2}
     wb::Array{Float64,1}
     f::Function
     df::Union{Function,Nothing}
     """
        DenseLayer(f,n,nₗ;w,wb,df)

     Instantiate a new DenseLayer

     # Positional arguments:
     * `nₗ`: Number of nodes of the previous layer
     * `n`:  Number of nodes
     # Keyword arguments:
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (n,nₗ)]
     * `wb`:  Initial weigths with respect to bias [default: Xavier initialisation, dims = (n)]
     * `f`:   Activation function [def: `identity`]
     * `df`:  Derivative of the activation function [default: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

     # Notes:
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n))`

     """
     function DenseLayer(nₗ,n;rng = Random.GLOBAL_RNG,w=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n,nₗ),wb=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n),f=identity,df=nothing)
         # To be sure w is a matrix and wb a column vector..
         w  = reshape(w,n,nₗ)
         wb = reshape(wb,n)
         return new(w,wb,f,df)
     end
end

function _zComp(layer::DenseLayer,x)
    w  = layer.w
    wb = layer.wb
    z  = zeros(eltype(x),size(w,1))
    @inbounds for n in axes(w,1)
        zn = zero(eltype(x))
        @simd for nl in axes(x,1)
            zn += w[n,nl] * x[nl]
        end
        zn   += wb[n]
        z[n]  = zn
    end
    return z
end

function forward(layer::DenseLayer,x)
  z =  _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) #layer.w * x + layer.wb # _zComp(layer,x) #   layer.w * x + layer.wb # testd @avx
  return layer.f.(z)
end

function backward(layer::DenseLayer,x,nextGradient)
   z = _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) # layer.w * x + layer.wb # _zComp(layer,x) # @avx layer.w * x + layer.wb               # tested @avx
   if layer.df != nothing
       dϵ_dz =  layer.df.(z) .* nextGradient # tested @avx
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dI =  layer.w' * dϵ_dz # @avx
end

function getParams(layer::DenseLayer)
  return Learnable((layer.w,layer.wb))
end

function getGradient(layer::DenseLayer,x,nextGradient)
   z      =  _zComp(layer,x) #@avx layer.w * x + layer.wb #  _zComp(layer,x) #layer.w * x + layer.wb # @avx
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dw  = dϵ_dz * x' # @avx
   dϵ_dwb = dϵ_dz
   return Learnable((dϵ_dw,dϵ_dwb))
end

function setParams!(layer::DenseLayer,w)
   layer.w  = w.data[1]
   layer.wb = w.data[2]
end
function size(layer::DenseLayer)
    return size(layer.w')
end

# ------------------------------------------------------------------------------
# DenseNoBiasLayer layer

"""
   DenseNoBiasLayer

Representation of a layer without bias in the network

# Fields:
* `w`:  Weigths matrix with respect to the input from previous layer or data (n x n pr. layer)
* `f`:  Activation function
* `df`: Derivative of the activation function
"""
mutable struct DenseNoBiasLayer <: AbstractLayer
     w::Array{Float64,2}
     f::Function
     df::Union{Function,Nothing}
     """
        DenseNoBiasLayer(f,nₗ,n;w,df)

     Instantiate a new DenseNoBiasLayer

        DenseNoBiasLayer(f,nₗ,n;w,df)

     Instantiate a new DenseNoBiasLayer

     # Positional arguments:
     * `nₗ`:  Number of nodes of the previous layer
     * `n`:   Number of nodes
     # Keyword arguments:
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (nₗ,n)]
     * `f`:   Activation function [def: `identity`]
     * `df`:  Derivative of the activation function [def: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]
     # Notes:
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ,n))`
     """
     function DenseNoBiasLayer(nₗ,n;rng = Random.GLOBAL_RNG,w=rand(rng,Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n,nₗ),f=identity,df=nothing)
         # To be sure w is a matrix and wb a column vector..
         w  = reshape(w,n,nₗ)
         return new(w,f,df)
     end
end

function _zComp(layer::DenseNoBiasLayer,x)
    w  = layer.w
    z  = zeros(eltype(x),size(w,1))
    @inbounds for n in axes(w,1)
        zn = zero(eltype(x))
        @simd for nl in axes(x,1)
            zn += w[n,nl] * x[nl]
        end
        z[n]  = zn
    end
    return z
end

function forward(layer::DenseNoBiasLayer,x)
  z = _zComp(layer,x)
  return layer.f.(z)
end

function backward(layer::DenseNoBiasLayer,x,nextGradient)
   z = _zComp(layer,x)
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dI = layer.w' * dϵ_dz
end

function getParams(layer::DenseNoBiasLayer)
  return Learnable((layer.w,))
end

function getGradient(layer::DenseNoBiasLayer,x,nextGradient)
   z      = _zComp(layer,x)
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dw  =  dϵ_dz * x'
   return Learnable((dϵ_dw,))
end

function setParams!(layer::DenseNoBiasLayer,w)
   layer.w = w.data[1]
end

function size(layer::DenseNoBiasLayer)
    return size(layer.w')
end

# ------------------------------------------------------------------------------
# VectorFunctionLayer layer

"""
   VectorFunctionLayer

Representation of a VectorFunction layer in the network. Vector
function layer expects a vector activation function, i.e. a function taking the
whole output of the previous layer an input rather than working on a single node
as "normal" activation functions would do.
Useful for example with the SoftMax function in classification or with the
`pool1D` function to implement a "pool" layer in 1 dimensions.
By default it is weightless, i.e. it doesn't apply any transformation to the
output coming from the previous layer except the activation function. However,
by passing the parameter `wsize` (a touple or array - tested only 1D) you can
pass the learnable parameter to the activation function too. It is your
responsability to be sure the activation function accept only X or also this
learnable array (as second argument).
The number of nodes in input must be set to the same as in the previous layer
(and if you are using this for classification, to the number of classes, i.e.
the _previous_ layer must be set equal to the number of classes in the
predictions).

# Fields:
* `w`:   Weigths (parameter) array passes as second argument to the activation
         function (if not empty)
* `nₗ`:  Number of nodes in input (i.e. length of previous layer)
* `n`:   Number of nodes in output (automatically inferred in the constructor)
* `f`:   Activation function (vector)
* `dfx`: Derivative of the (vector) activation function with respect to the
         layer inputs (x)
* `dfw`: Derivative of the (vector) activation function with respect to the
         optional learnable weigths (w)

# Notes:
* The output `size` of this layer is given by the size of the output function,
that not necessarily is the same as the previous layers.
"""
mutable struct VectorFunctionLayer{N} <: AbstractLayer
     w::Array{Float64,N}
     nₗ::Int64
     n::Int64
     f::Function
     dfx::Union{Function,Nothing}
     dfw::Union{Function,Nothing}
     """
        VectorFunctionLayer(n;rng,wsize,w,f,dfx,dfw,dummyDataToTestOutputSize)

     Instantiate a new VectorFunctionLayer

     # Positional arguments:
     * `nₗ`: Number of nodes (must be same as in the previous layer)
     # Keyword arguments:
     * `wsize`: A tuple or array specifying the size (number of elements) of the
       learnable parameter [def: empty array]
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (nₗ,n)]
     * `f`:  Activation function [def: `softMax`]
     * `dfx`: Derivative of the activation function with respect to the data
              [default: `nothing` (i.e. use AD)]
     * `dfw`: Derivative of the activation function with respect to the
              learnable parameter [default: `nothing` (i.e. use AD)]
     * `dummyDataToTestOutputSize`: Dummy data to test the output size [def:
     `ones(nₗ)`]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]
     # Notes:
     - If the derivative is provided, it should return the gradient as a (n,n) matrix (i.e. the Jacobian)
     - To avoid recomputing the activation function just to determine its output size,
       we compute the output size once here in the layer constructor by calling the
       activation function with `dummyDataToTestOutputSize`. Feel free to change
       it if it doesn't match with the activation function you are setting
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(sum(wsize...)),sqrt(6)/sqrt(sum(wsize...))))`
     """
     function VectorFunctionLayer(nₗ;rng = Random.GLOBAL_RNG,wsize=Int64[],w=rand(rng,Uniform(-sqrt(6)/sqrt(sum(wsize)),sqrt(6)/sqrt(sum(wsize))),Tuple(wsize)), f=softMax,dfx=nothing,dfw=nothing,dummyDataToTestOutputSize=ones(nₗ))
        nw = length(wsize)
        if nw ==0
          n  = length(f(dummyDataToTestOutputSize))
        else
         n  = length(f(dummyDataToTestOutputSize,w))
        end
        return new{nw}(w,nₗ,n,f,dfx,dfw)
     end
end

function forward(layer::VectorFunctionLayer{N},x) where {N}
  return N == 0 ? layer.f(x) : layer.f(x,layer.w)
end

function backward(layer::VectorFunctionLayer{N},x,nextGradient) where N
   if N == 0
      if layer.dfx != nothing
         dϵ_dI = layer.dfx(x)' * nextGradient
      else  # using AD
         j = autoJacobian(layer.f,x; nY=layer.n)
         dϵ_dI = j' * nextGradient
      end
      return dϵ_dI
   else
      if layer.dfx != nothing
         dϵ_dI = layer.dfx(x,layer.w)' * nextGradient
      else  # using AD
        tempfunction(x) = layer.f(x,layer.w)
        nYl = layer.n
        j = autoJacobian(tempfunction,x; nY=nYl)
        dϵ_dI = j' * nextGradient
      end
     return dϵ_dI
   end
end

function getParams(layer::VectorFunctionLayer{N}) where {N}
   return N == 0 ? Learnable(()) : Learnable((layer.w,))
end

function getGradient(layer::VectorFunctionLayer{N},x,nextGradient) where {N}
   if N == 0
     return Learnable(()) # parameterless layer
   else
      if layer.dfw != nothing
         dϵ_dw = layer.dfw(x,layer.w)' * nextGradient
      else  # using AD
        j = autoJacobian(wt -> layer.f(x,wt),layer.w; nY=layer.n)
        dϵ_dw = j' * nextGradient
      end
     return Learnable((dϵ_dw,))
   end
end

function setParams!(layer::VectorFunctionLayer{N},w) where {N}
   if N > 0
      layer.w = w.data[1]
   end
   return nothing
end
function size(layer::VectorFunctionLayer{N}) where {N}
    # Output size for the VectorFunctionLayer is given by its activation function
    # We test its length with dummy values
    return (layer.nₗ,layer.n)
end

# ------------------------------------------------------------------------------
# ScalarFunctionLayer layer

"""
   ScalarFunctionLayer

Representation of a ScalarFunction layer in the network.
ScalarFunctionLayer applies the activation function directly to the output of
the previous layer (i.e., without passing for a weigth matrix), but using an
optional learnable parameter (an array) used as second argument, similarly to
[`VectorFunctionLayer`(@ref).
Differently from `VectorFunctionLayer`, the function is applied scalarwise to
each node.

The number of nodes in input must be set to the same as in the previous layer

# Fields:
* `w`:   Weigths (parameter) array passes as second argument to the activation
         function (if not empty)
* `n`:   Number of nodes in output (≡ number of nodes in input )
* `f`:   Activation function (vector)
* `dfx`: Derivative of the (vector) activation function with respect to the
         layer inputs (x)
* `dfw`: Derivative of the (vector) activation function with respect to the
         optional learnable weigths (w)

# Notes:
* The output `size` of this layer is the same as those of the previous layers.
"""
mutable struct ScalarFunctionLayer{N} <: AbstractLayer
     w::Array{Float64,N}
     n::Int64
     f::Function
     dfx::Union{Function,Nothing}
     dfw::Union{Function,Nothing}
     """
        ScalarFunctionLayer(n;rng,wsize,w,f,dfx,dfw)

     Instantiate a new ScalarFunctionLayer

     # Positional arguments:
     * `nₗ`: Number of nodes (must be same as in the previous layer)
     # Keyword arguments:
     * `wsize`: A tuple or array specifying the size (number of elements) of the
       learnable parameter [def: empty array]
     * `w`:   Initial weigths with respect to input [default: Xavier initialisation, dims = (nₗ,n)]
     * `f`:  Activation function [def: `softMax`]
     * `dfx`: Derivative of the activation function with respect to the data
              [default: `nothing` (i.e. use AD)]
     * `dfw`: Derivative of the activation function with respect to the
              learnable parameter [default: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]
     # Notes:
     - If the derivative is provided, it should return the gradient as a (n,n) matrix (i.e. the Jacobian)
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(sum(wsize...)),sqrt(6)/sqrt(sum(wsize...))))`
     """
     function ScalarFunctionLayer(nₗ;rng = Random.GLOBAL_RNG,wsize=Int64[],w=rand(rng,Uniform(-sqrt(6)/sqrt(sum(wsize)),sqrt(6)/sqrt(sum(wsize))),Tuple(wsize)), f=softMax,dfx=nothing,dfw=nothing)
        nw = length(wsize)
        return new{nw}(w,nₗ,f,dfx,dfw)
     end
end

function forward(layer::ScalarFunctionLayer{N},x) where {N}
  return N == 0 ? layer.f.(x) : layer.f.(x,layer.w)
end

function backward(layer::ScalarFunctionLayer{N},x,nextGradient) where N
   if N == 0
      if layer.dfx != nothing
         dϵ_dI = layer.dfx.(x) .* nextGradient
      else  # using AD
         dϵ_dI = layer.f'.(x) .* nextGradient
      end
      return dϵ_dI
   else
      if layer.dfx != nothing
         dϵ_dI = layer.dfx.(x,Ref(layer.w)) .* nextGradient
      else  # using AD
        #tempfunction(x) = layer.f.(x,Ref(layer.w))
        df_dx = [gradient(xt -> layer.f(xt,layer.w),xi)[1] for xi in x]
        dϵ_dI = df_dx .* nextGradient
      end
     return dϵ_dI
   end
end

function getParams(layer::ScalarFunctionLayer{N}) where {N}
   return N == 0 ? Learnable(()) : Learnable((layer.w,))
end

function getGradient(layer::ScalarFunctionLayer{N},x,nextGradient) where {N}
   if N == 0
     return Learnable(()) # parameterless layer
   else
      if layer.dfw != nothing
         dϵ_dw = [layer.dfw(xi,wj) for wj in layer.w, xi in x] * nextGradient
      else  # using AD
        tempfunction(w) = [layer.f(xi,w) for xi in x]
        j = autoJacobian(tempfunction,layer.w; nY=layer.n)
        dϵ_dw = j' * nextGradient
      end
     return Learnable((dϵ_dw,))
   end
end

function setParams!(layer::ScalarFunctionLayer{N},w) where {N}
   if N > 0
      layer.w = w.data[1]
   end
   return nothing
end
function size(layer::ScalarFunctionLayer{N}) where {N}
    # Output size for the ScalarFunctionLayer is given by its activation function
    # We test its length with dummy values
    return (layer.n,layer.n)
end

# ------------------------------------------------------------------------------
# RNN layer

"""
   RNNLayer

Representation of a layer in the network

# Fields:
* `wx`: Weigths matrix with respect to the input from data (n by n_input)
* `ws`: Weigths matrix with respect to the layer state (n x n )
* `wb`: Biases (n)
* `f`:  Activation function
* `df`: Derivative of the activation function
* `s` : State
"""
mutable struct RNNLayer <: RecursiveLayer
     wx::Array{Float64,2}
     ws::Array{Float64,2}
     wb::Array{Float64,1}
     s::Array{Float64,1}
     f::Function
     df::Union{Function,Nothing}
     """
        RNNLayer(nₗ,n;f,wx,ws,wb,df)

     Instantiate a new RNNLayer

     # Positional arguments:
     * `nₗ`: Number of nodes of the input
     * `n`:  Number of nodes of the state (and the output)
     # Keyword arguments:
     * `wx`:  Initial weigths with respect to input [default: Xavier initialisation, dims = (n,nₗ)]
     * `ws`:  Initial weigths with respect to input [default: Xavier initialisation, dims = (n,n)]
     * `wb`:  Initial weigths with respect to bias [default: Xavier initialisation, dims = (n)]
     * `s`:   Initial states [def: zeros(n)]
     * `f`:   Activation function [def: `relu`]
     * `df`:  Derivative of the activation function [default: `nothing` (i.e. use AD)]
     * `rng`: Random Number Generator (see [`FIXEDSEED`](@ref)) [deafult: `Random.GLOBAL_RNG`]

     # Notes:
     - Xavier initialization = `rand(Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n))`

     """
     function RNNLayer(nₗ,n;rng = Random.GLOBAL_RNG,wx=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n,nₗ),ws=rand(rng, Uniform(-sqrt(6)/sqrt(n+n),sqrt(6)/sqrt(n+n)),n,n),wb=rand(rng, Uniform(-sqrt(6)/sqrt(nₗ+n),sqrt(6)/sqrt(nₗ+n)),n),s=zeros(n),f=relu,df=nothing)
         # To be sure w is a matrix and wb a column vector..
         wx  = reshape(w,n,nₗ)
         ws  = reshape(w,n,n)
         wb  = reshape(wb,n)
         s   = reshape(s,n)
         return new(w,wb,s,f,df)
     end
end
#=
function _zComp(layer::RNNLayer,x)
    w  = layer.w
    wb = layer.wb
    z  = zeros(eltype(x),size(w,1))
    @inbounds for n in axes(w,1)
        zn = zero(eltype(x))
        @simd for nl in axes(x,1)
            zn += w[n,nl] * x[nl]
        end
        zn   += wb[n]
        z[n]  = zn
    end
    return z
end
=#

function forward(layer::RNNLayer,x)
  #z =  _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) #layer.w * x + layer.wb # _zComp(layer,x) #   layer.w * x + layer.wb # testd @avx
  z =  layer.wb + layer.wx * x + layer.ws * s
  return layer.f.(z)
end

function backward(layer::RNNLayer,x,nextGradient) #TODO
   z = _zComp(layer,x) #@avx layer.w * x + layer.wb #_zComp(layer,x) # layer.w * x + layer.wb # _zComp(layer,x) # @avx layer.w * x + layer.wb               # tested @avx
   if layer.df != nothing
       dϵ_dz =  layer.df.(z) .* nextGradient # tested @avx
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dI =  layer.w' * dϵ_dz # @avx
end

function getParams(layer::RNNLayer)
  return Learnable((layer.wb,layer.wx,layer.ws))
end

function getGradient(layer::RNNLayer,x,nextGradient) #TODO
   z      =  _zComp(layer,x) #@avx layer.w * x + layer.wb #  _zComp(layer,x) #layer.w * x + layer.wb # @avx
   if layer.df != nothing
       dϵ_dz = layer.df.(z) .* nextGradient
    else
       dϵ_dz = layer.f'.(z) .* nextGradient # using AD
    end
   dϵ_dw  = dϵ_dz * x' # @avx
   dϵ_dwb = dϵ_dz
   return Learnable((dϵ_dw,dϵ_dwb))
end

function setParams!(layer::RNNLayer,w)
   layer.wb = w.data[1]
   layer.wx = w.data[2]
   layer.ws = w.data[3]
end
function size(layer::RNNLayer)
    return size(layer.w')
end


"""
    SimpleANN <: AbstractNeuralNetwork{Int,Int,1}

A simple, 1-layer, feed-forward artificial neural network with perceptron nodes.

    SimpleANN(ws::AbstractArray{Float64,2}, ts::AbstractVector{Float64})

Create a SimpleANN with weights `ws` and thresholds `ts`.

    SimpleANN(n::Int,m::Int) = SimpleANN(zeros(m,n), zeros(m))

Create a SimpleANN with `n` inputs and `m` outputs. All weights and thresholds are
initialized to zero.
"""
struct SimpleANN <: AbstractNeuralNetwork{Int,Int,1}
    insize :: Int
    outsize :: Int
    weights :: Array{Float64,2}
    thresholds :: Vector{Float64}
    function SimpleANN(ws::AbstractArray{Float64,2}, ts::AbstractVector{Float64})
        if any(size(ws).==0)
            throw(ArgumentError("weights matrix; invalid size"))
        elseif length(ts) != size(ws,1)
            throw(DimensionMismatch("weights and thresholds"))
        end
        m, n = size(ws)
        new(n, m, ws, ts)
    end
end

SimpleANN(n::Int,m::Int) = SimpleANN(zeros(m,n), zeros(m))

insize(net::SimpleANN) = (net.insize,)
outsize(net::SimpleANN) = (net.outsize,)

function fire!(net::SimpleANN, input::AbstractVector{Int},
               output::AbstractVector{Int})
    if net.insize != length(input)
        throw(DimensionMismatch("invalid input"))
    elseif net.outsize != length(output)
        throw(DimensionMismatch("invalid output"))
    end
    temp = net.weights * input - net.thresholds
    for i in 1:net.outsize
        output[i] = (temp[i] > 0.0)
    end
    output
end


"""
    VectorSpaceType

Abstract type for tangent spaces, cotangent spaces, their tensor products,
exterior products, etc.

Every vector space `fiber` is supposed to provide:
* a method of constructing vectors,
* basic operations: addition, subtraction, multiplication by a scalar
  and negation (unary minus),
* [`zero_vector!(fiber, X, p)`](@ref) to construct zero vectors at point `p`,
* `allocate(X)` and `allocate(X, T)` for vector `X` and type `T`,
* `copyto!(X, Y)` for vectors `X` and `Y`,
* `number_eltype(v)` for vector `v`,
* [`vector_space_dimension(::VectorBundleFibers{<:typeof(fiber)}) where fiber`](@ref).

Optionally:
* inner product via `inner` (used to provide Riemannian metric on vector
  bundles),
* [`flat`](@ref) and [`sharp`](@ref),
* `norm` (by default uses `inner`),
* [`project_vector`](@ref) (for embedded vector spaces),
* [`representation_size`](@ref) (if support for [`ProductArray`](@ref) is desired),
* broadcasting for basic operations.
"""
abstract type VectorSpaceType end

struct TangentSpaceType <: VectorSpaceType end

struct CotangentSpaceType <: VectorSpaceType end

TCoTSpaceType = Union{TangentSpaceType,CotangentSpaceType}

const TangentSpace = TangentSpaceType()
const CotangentSpace = CotangentSpaceType()

"""
    TensorProductType(spaces::VectorSpaceType...)

Vector space type corresponding to the tensor product of given vector space
types.
"""
struct TensorProductType{TS<:Tuple} <: VectorSpaceType
    spaces::TS
end

TensorProductType(spaces::VectorSpaceType...) = TensorProductType{typeof(spaces)}(spaces)

"""
    VectorBundleFibers(fiber::VectorSpaceType, M::Manifold)

Type representing a family of vector spaces (fibers) of a vector bundle over `M`
with vector spaces of type `fiber`. In contrast with `VectorBundle`, operations
on `VectorBundleFibers` expect point-like and vector-like parts to be
passed separately instead of being bundled together. It can be thought of
as a representation of vector spaces from a vector bundle but without
storing the point at which a vector space is attached (which is specified
separately in various functions).
"""
struct VectorBundleFibers{TVS<:VectorSpaceType,TM<:Manifold}
    fiber::TVS
    manifold::TM
end

const TangentBundleFibers{M} = VectorBundleFibers{TangentSpaceType,M} where {M<:Manifold}

TangentBundleFibers(M::Manifold) = VectorBundleFibers(TangentSpace, M)

const CotangentBundleFibers{M} =
    VectorBundleFibers{CotangentSpaceType,M} where {M<:Manifold}

CotangentBundleFibers(M::Manifold) = VectorBundleFibers(CotangentSpace, M)

"""
    VectorSpaceAtPoint(fiber::VectorBundleFibers, p)

A vector space (fiber type `fiber` of a vector bundle) at point `p` from
the manifold `fiber.manifold`.
"""
struct VectorSpaceAtPoint{TFiber<:VectorBundleFibers,TX}
    fiber::TFiber
    point::TX
end

"""
    TangentSpaceAtPoint(M::Manifold, p)

Return an object of type [`VectorSpaceAtPoint`](@ref) representing tangent
space at `p`.
"""
TangentSpaceAtPoint(M::Manifold, p) = VectorSpaceAtPoint(TangentBundleFibers(M), p)

"""
    CotangentSpaceAtPoint(M::Manifold, p)

Return an object of type [`VectorSpaceAtPoint`](@ref) representing cotangent
space at `p`.
"""
CotangentSpaceAtPoint(M::Manifold, p) = VectorSpaceAtPoint(CotangentBundleFibers(M), p)

"""
    VectorBundle(M::Manifold, type::VectorSpaceType)

Vector bundle on manifold `M` of type `type`.
"""
struct VectorBundle{TVS<:VectorSpaceType,TM<:Manifold} <: Manifold
    type::TVS
    manifold::TM
    fiber::VectorBundleFibers{TVS,TM}
end

function VectorBundle(fiber::TVS, M::TM) where {TVS<:VectorSpaceType,TM<:Manifold}
    return VectorBundle{TVS,TM}(fiber, M, VectorBundleFibers(fiber, M))
end

const TangentBundle{M} = VectorBundle{TangentSpaceType,M} where {M<:Manifold}

TangentBundle(M::Manifold) = VectorBundle(TangentSpace, M)

const CotangentBundle{M} = VectorBundle{CotangentSpaceType,M} where {M<:Manifold}

CotangentBundle(M::Manifold) = VectorBundle(CotangentSpace, M)

"""
    FVector(type::VectorSpaceType, data)

Decorator indicating that the vector `data` is from a fiber of a vector bundle
of type `type`.
"""
struct FVector{TType<:VectorSpaceType,TData}
    type::TType
    data::TData
end

const TFVector = FVector{TangentSpaceType}
const CoTFVector = FVector{CotangentSpaceType}

struct VectorBundleBasisData{BBasis<:CachedBasis, TBasis<:CachedBasis}
    base_basis::BBasis
    vec_basis::TBasis
end

(+)(X::FVector, Y::FVector) = FVector(X.type, X.data + Y.data)

(-)(X::FVector, Y::FVector) = FVector(X.type, X.data - Y.data)
(-)(X::FVector) = FVector(X.type, -X.data)

(*)(a::Number, X::FVector) = FVector(X.type, a * X.data)

function copyto!(X::FVector, Y::FVector)
    copyto!(X.data, Y.data)
    return X
end

base_manifold(B::VectorBundleFibers) = base_manifold(B.manifold)
base_manifold(B::VectorSpaceAtPoint) = base_manifold(B.fiber)
base_manifold(B::VectorBundle) = base_manifold(B.manifold)

"""
    bundle_projection(B::VectorBundle, x::ProductRepr)

Projection of point `p` from the bundle `M` to the base manifold.
Returns the point on the base manifold `B.manifold` at which the vector part
of `p` is attached.
"""
bundle_projection(B::VectorBundle, p) = submanifold_component(B.manifold, p, Val(1))

"""
    distance(B::VectorBundleFibers, p, X, Y)

Distance between vectors `X` and `Y` from the vector space at point `p`
from the manifold `B.manifold`, that is the base manifold of `M`.
"""
distance(B::VectorBundleFibers, p, X, Y) = norm(B, p, X - Y)
@doc raw"""
    distance(B::VectorBundle, p, q)

Distance between points $x$ and $y$ from the
vector bundle `B` over manifold `B.fiber` (denoted $\mathcal M$).

Notation:
  * The point $p = (x_p, V_p)$ where $x_p  ∈ \mathcal M$ and $V_p$ belongs to the
    fiber $F=π^{-1}(\{x_p\})$ of the vector bundle $B$ where $π$ is the
    canonical projection of that vector bundle $B$.
    Similarly, $q = (x_q, V_q)$.

The distance is calculated as

$d_B(x, y) = \sqrt{d_M(x_p, x_q)^2 + d_F(V_p, V_{q←p})^2}$

where $d_\mathcal M$ is the distance on manifold $\mathcal M$, $d_F$ is the distance
between two vectors from the fiber $F$ and $V_{q←p}$ is the result
of parallel transport of vector $V_q$ to point $x_p$. The default
behavior of [`vector_transport_to`](@ref) is used to compute the vector
transport.
"""
function distance(B::VectorBundle, p, q)
    xp, Vp = submanifold_components(B.manifold, p)
    xq, Vq = submanifold_components(B.manifold, q)
    dist_man = distance(B.manifold, xp, xq)
    vy_x = vector_transport_to(B.manifold, xq, Vq, xp)
    dist_vec = distance(B.fiber, xp, Vp, vy_x)
    return sqrt(dist_man^2 + dist_vec^2)
end

function number_eltype(::Type{FVector{TType,TData}}) where {TType<:VectorSpaceType,TData}
    return number_eltype(TData)
end
number_eltype(v::FVector) = number_eltype(v.data)

@doc raw"""
    exp(B::VectorBundle, p, X)

Exponential map of tangent vector $X$ at point $p$ from
vector bundle `B` over manifold `B.fiber` (denoted $\mathcal M$).

Notation:
  * The point $p = (x_p, V_p)$ where $x_p ∈ \mathcal M$ and $V_p$ belongs to the
    fiber $F=π^{-1}(\{x_p\})$ of the vector bundle $B$ where $π$ is the
    canonical projection of that vector bundle $B$.
  * The tangent vector $X = (V_{X,M}, V_{X,F}) ∈ T_pB$ where
    $V_{X,M}$ is a tangent vector from the tangent space $T_{x_p}\mathcal M$ and
    $V_{X,F}$ is a tangent vector from the tangent space $T_{V_p}F$ (isomorphic to $F$).

The exponential map is calculated as

$\exp_p(X) = (\exp_{x_p}(V_{X,M}), V_{\exp})$

where $V_{\exp}$ is the result of vector transport of $V_p + V_{X,F}$
to the point $\exp_{x_p}(V_{X,M})$.
The sum $V_p + V_{X,F}$ corresponds to the exponential map in the vector space $F$.
"""
exp(::VectorBundle, ::Any)

function exp!(B::VectorBundle, q, p, X)
    xp, Xp = submanifold_components(B.manifold, p)
    xq, Xq = submanifold_components(B.manifold, q)
    VXM, VXF = submanifold_components(B.manifold, X)
    exp!(B.manifold, xq, xp, VXM)
    vector_transport_to!(B.manifold, Xq, xp, Xp + VXF, xq)
    return q
end

@doc raw"""
    flat(M::Manifold, p, X::FVector)

Compute the flat isomorphism (one of the musical isomorphisms) of tangent vector `X`
from the vector space of type `M` at point `p` from the underlying [`Manifold`](@ref).

The function can be used for example to transform vectors
from the tangent bundle to vectors from the cotangent bundle
$♭ : T\mathcal M → T^{*}\mathcal M$
"""
function flat(M::Manifold, p, X::FVector)
    ξ = allocate_result(M, flat, X, p)
    return flat!(M, ξ, p, X)
end

function flat!(M::Manifold, ξ::FVector, p, X::FVector)
    error(
        "flat! not implemented for vector bundle fibers space " *
        "of type $(typeof(M)), vector of type $(typeof(ξ)), point of " *
        "type $(typeof(p)) and vector of type $(typeof(X)).",
    )
end

@decorator_transparent_signature flat!(
    M::AbstractDecoratorManifold,
    ξ::CoTFVector,
    p,
    X::TFVector,
)

function get_basis(M::VectorBundle, p, B::AbstractBasis)
    xp1 = submanifold_component(p, Val(1))
    base_basis = get_basis(M.manifold, xp1, B)
    vec_basis = get_basis(M.fiber, xp1, B)
    return CachedBasis(B, VectorBundleBasisData(base_basis, vec_basis))
end
function get_basis(M::VectorBundle, p, B::CachedBasis)
    return invoke(get_basis, Tuple{Manifold, Any, CachedBasis}, M, p, B)
end

function get_basis(M::VectorBundle, p, B::DiagonalizingOrthonormalBasis)
    xp1 = submanifold_component(p, Val(1))
    bv1 = DiagonalizingOrthonormalBasis(submanifold_component(B.frame_direction, Val(1)))
    b1 = get_basis(M.manifold, xp1, bv1)
    bv2 = DiagonalizingOrthonormalBasis(submanifold_component(B.frame_direction, Val(2)))
    b2 = get_basis(M.fiber, xp1, bv2)
    return CachedBasis(B, VectorBundleBasisData(b1, b2))
end

for BT in [
    DefaultOrthonormalBasis,
    ProjectedOrthonormalBasis{:gram_schmidt,ℝ},
    ProjectedOrthonormalBasis{:svd,ℝ},
]
    eval(quote
        @invoke_maker 3 AbstractBasis get_basis(M::VectorBundle, p, B::$BT)
    end)
end
function get_basis(M::TangentBundleFibers, p, B::AbstractBasis)
    return get_basis(M.manifold, p, B)
end

function get_coordinates!(M::VectorBundle, Y, p, X, B::AbstractBasis)
    px, Vx = submanifold_components(M.manifold, p)
    VXM, VXF = submanifold_components(M.manifold, X)
    n = manifold_dimension(M.manifold)
    get_coordinates!(M.manifold, view(Y, 1:n), px, VXM, B)
    get_coordinates!(M.fiber, view(Y, n+1:length(Y)), px, VXF, B)
    return Y
end
function get_coordinates!(
    M::VectorBundle,
    Y,
    p,
    X,
    B::CachedBasis{<:AbstractBasis,<:VectorBundleBasisData},
) where {N}
     px, Vx = submanifold_components(M.manifold, p)
     VXM, VXF = submanifold_components(M.manifold, X)
     n = manifold_dimension(M.manifold)
     get_coordinates!(M.manifold, view(Y, 1:n), px, VXM, B.data.base_basis)
     get_coordinates!(M.fiber, view(Y, n+1:length(Y)), px, VXF, B.data.vec_basis)
     return Y
end
for BT in [
    DefaultBasis,
    DefaultOrthogonalBasis,
    DefaultOrthonormalBasis,
    ProjectedOrthonormalBasis{:gram_schmidt,ℝ},
    ProjectedOrthonormalBasis{:svd,ℝ},
    VeeOrthogonalBasis,
]
    eval(quote
        @invoke_maker 5 AbstractBasis get_coordinates!(M::VectorBundle, Y, p, X, B::$BT)
    end)
end
for BT in [
    CachedBasis{<:AbstractBasis{ℝ},<:VectorBundleBasisData},
    CachedBasis{<:ManifoldsBase.AbstractOrthogonalBasis{ℝ},<:VectorBundleBasisData},
    CachedBasis{<:ManifoldsBase.AbstractOrthonormalBasis{ℝ},<:VectorBundleBasisData},
    CachedBasis{<:AbstractBasis{ℂ},<:VectorBundleBasisData},
]
    eval(quote
        @invoke_maker 5 CachedBasis{<:AbstractBasis,<:VectorBundleBasisData} get_coordinates!(
                M::VectorBundle,
                Y,
                p,
                X,
                B::$BT,
            )
    end)
end

function get_coordinates!(M::TangentBundleFibers, Y, p, X, B::ManifoldsBase.all_uncached_bases) where {N}
    return get_coordinates!(M.manifold, Y, p, X, B)
end

function get_vector!(M::VectorBundle, Y, p, X, B::DefaultOrthonormalBasis)
    n = manifold_dimension(M.manifold)
    xp1 = submanifold_component(p, Val(1))
    get_vector!(M.manifold, submanifold_component(Y, Val(1)), xp1, X[1:n], B)
    get_vector!(M.fiber, submanifold_component(Y, Val(2)), xp1, X[n+1:end], B)
    return Y
end
function get_vector!(
    M::VectorBundle,
    Y,
    p,
    X,
    B::CachedBasis{<:AbstractBasis,<:VectorBundleBasisData},
) where {N}
    n = manifold_dimension(M.manifold)
    xp1 = submanifold_component(p, Val(1))
    get_vector!(
        M.manifold,
        submanifold_component(Y, Val(1)),
        xp1,
        X[1:n],
        B.data.base_basis,
    )
    get_vector!(
        M.fiber,
        submanifold_component(Y, Val(2)),
        xp1,
        X[n+1:end],
        B.data.vec_basis,
    )
    return Y
end
function get_vector!(M::TangentBundleFibers, Y, p, X, B::ManifoldsBase.all_uncached_bases) where {N}
    return get_vector!(M.manifold, Y, p, X, B)
end

function get_vectors(
    M::VectorBundle,
    p,
    B::CachedBasis{<:AbstractBasis,<:VectorBundleBasisData},
)
    xp1 = submanifold_component(p, Val(1))
    zero_m = zero_tangent_vector(M.manifold, xp1)
    zero_f = zero_vector(M.fiber, xp1)
    vs = typeof(ProductRepr(zero_m, zero_f))[]
    for bv in get_vectors(M.manifold, xp1, B.data.base_basis)
        push!(vs, ProductRepr(bv, zero_f))
    end
    for bv in get_vectors(M.fiber, xp1, B.data.vec_basis)
        push!(vs, ProductRepr(zero_m, bv))
    end
    return vs
end

function get_vectors(M::VectorBundleFibers, p, B::CachedBasis)
    return get_vectors(M.manifold, p, B)
end
Base.@propagate_inbounds getindex(x::FVector, i) = getindex(x.data, i)

"""
    inner(B::VectorBundleFibers, p, X, Y)

Inner product of vectors `X` and `Y` from the vector space of type `B.fiber`
at point `p` from manifold `B.manifold`.
"""
function inner(B::VectorBundleFibers, p, X, Y)
    error(
        "inner not defined for vector space family of type $(typeof(B)), " *
        "point of type $(typeof(p)) and " *
        "vectors of types $(typeof(X)) and $(typeof(Y)).",
    )
end
inner(B::VectorBundleFibers{<:TangentSpaceType}, p, X, Y) = inner(B.manifold, p, X, Y)
function inner(B::VectorBundleFibers{<:CotangentSpaceType}, p, X, Y)
    return inner(
        B.manifold,
        p,
        sharp(B.manifold, p, FVector(CotangentSpace, X)).data,
        sharp(B.manifold, p, FVector(CotangentSpace, Y)).data,
    )
end
@doc raw"""
    inner(B::VectorBundle, p, X, Y)

Inner product of tangent vectors `X` and `Y` at point `p` from the
vector bundle `B` over manifold `B.fiber` (denoted $\mathcal M$).

Notation:
  * The point $p = (x_p, V_p)$ where $x_p ∈ \mathcal M$ and $V_p$ belongs to the
    fiber $F=π^{-1}(\{x_p\})$ of the vector bundle $B$ where $π$ is the
    canonical projection of that vector bundle $B$.
  * The tangent vector $v = (V_{X,M}, V_{X,F}) ∈ T_{x}B$ where
    $V_{X,M}$ is a tangent vector from the tangent space $T_{x_p}\mathcal M$ and
    $V_{X,F}$ is a tangent vector from the tangent space $T_{V_p}F$ (isomorphic to $F$).
    Similarly for the other tangent vector $w = (V_{Y,M}, V_{Y,F}) ∈ T_{x}B$.

The inner product is calculated as

$⟨X, Y⟩_p = ⟨V_{X,M}, V_{Y,M}⟩_{x_p} + ⟨V_{X,F}, V_{Y,F}⟩_{V_p}.$
"""
function inner(B::VectorBundle, p, X, Y)
    px, Vx = submanifold_components(B.manifold, p)
    VXM, VXF = submanifold_components(B.manifold, X)
    VYM, VYF = submanifold_components(B.manifold, Y)
    return inner(B.manifold, px, VXM, VYM) + inner(B.fiber, Vx, VXF, VYF)
end

function isapprox(B::VectorBundle, p, q; kwargs...)
    xp, Vp = submanifold_components(B.manifold, p)
    xq, Vq = submanifold_components(B.manifold, q)
    return isapprox(B.manifold, xp, xq; kwargs...) && isapprox(Vp, Vq; kwargs...)
end
function isapprox(B::VectorBundle, p, X, Y; kwargs...)
    px, Vx = submanifold_components(B.manifold, p)
    VXM, VXF = submanifold_components(B.manifold, X)
    VYM, VYF = submanifold_components(B.manifold, Y)
    return isapprox(B.manifold, VXM, VYM; kwargs...) &&
           isapprox(B.manifold, px, VXF, VYF; kwargs...)
end

@doc raw"""
    log(B::VectorBundle, p, q)

Logarithmic map of the point `y` at point `p` from
vector bundle `B` over manifold `B.fiber` (denoted $\mathcal M$).

Notation:
  * The point $p = (x_p, V_p)$ where $x_p ∈ \mathcal M$ and $V_p$ belongs to the
    fiber $F=π^{-1}(\{x_p\})$ of the vector bundle $B$ where $π$ is the
    canonical projection of that vector bundle $B$.
    Similarly, $q = (x_q, V_q)$.

The logarithmic map is calculated as

$\log_p q = (\log_{x_p}(x_q), V_{\log} - V_p)$

where $V_{\log}$ is the result of vector transport of $V_q$ to the point $x_p$.
The difference $V_{\log} - V_p$ corresponds to the logarithmic map in the vector space $F$.
"""
log(::VectorBundle, ::Any...)

function log!(B::VectorBundle, X, p, q)
    px, Vx = submanifold_components(B.manifold, p)
    py, Vy = submanifold_components(B.manifold, q)
    VXM, VXF = submanifold_components(B.manifold, X)
    log!(B.manifold, VXM, px, py)
    vector_transport_to!(B.manifold, VXF, py, Vy, px)
    copyto!(VXF, VXF - Vx)
    return X
end

function manifold_dimension(B::VectorBundle)
    return manifold_dimension(B.manifold) + vector_space_dimension(B.fiber)
end

"""
    norm(B::VectorBundleFibers, p, q)

Norm of the vector `X` from the vector space of type `B.fiber`
at point `p` from manifold `B.manifold`.
"""
norm(B::VectorBundleFibers, p, X) = sqrt(inner(B, p, X, X))
norm(B::VectorBundleFibers{<:TangentSpaceType}, p, X) = norm(B.manifold, p, X)

@doc raw"""
    project(B::VectorBundle, p)

Project the point `p` from the ambient space of the vector bundle `B`
over manifold `B.fiber` (denoted $\mathcal M$) to the vector bundle.

Notation:
  * The point $p = (x_p, V_p)$ where $x_p$ belongs to the ambient space of $\mathcal M$
    and $V_p$ belongs to the ambient space of the
    fiber $F=π^{-1}(\{x_p\})$ of the vector bundle $B$ where $π$ is the
    canonical projection of that vector bundle $B$.

The projection is calculated by projecting the point $x_p$ to the manifold $\mathcal M$
and then projecting the vector $V_p$ to the tangent space $T_{x_p}\mathcal M$.
"""
project(::VectorBundle, ::Any)

function project!(B::VectorBundle, q, p)
    px, pVx = submanifold_components(B.manifold, p)
    qx, qVx = submanifold_components(B.manifold, q)
    project!(B.manifold, qx, px)
    project!(B.manifold, qVx, qx, pVx)
    return q
end

@doc raw"""
    project(B::VectorBundle, p, X)

Project the element `X` of the ambient space of the tangent space $T_p B$
to the tangent space $T_p B$.

Notation:
  * The point $p = (x_p, V_p)$ where $x_p ∈ \mathcal M$ and $V_p$ belongs to the
    fiber $F=π^{-1}(\{x_p\})$ of the vector bundle $B$ where $π$ is the
    canonical projection of that vector bundle $B$.
  * The vector $x = (V_{X,M}, V_{X,F})$ where $x_p$ belongs to the ambient space of $T_{x_p}\mathcal M$
    and $V_{X,F}$ belongs to the ambient space of the
    fiber $F=π^{-1}(\{x_p\})$ of the vector bundle $B$ where $π$ is the
    canonical projection of that vector bundle $B$.

The projection is calculated by projecting $V_{X,M}$ to tangent space $T_{x_p}\mathcal M$
and then projecting the vector $V_{X,F}$ to the fiber $F$.
"""
project(::VectorBundle, ::Any, ::Any)

function project!(B::VectorBundle, Y, p, X)
    px, Vx = submanifold_components(B.manifold, p)
    VXM, VXF = submanifold_components(B.manifold, X)
    VYM, VYF = submanifold_components(B.manifold, Y)
    project!(B.manifold, VYM, px, VXM)
    project!(B.manifold, VYF, px, VXF)
    return Y
end

"""
    project_vector(B::VectorBundleFibers, p, X)

Project vector `X` from the vector space of type `B.fiber` at point `p`.
"""
function project_vector(B::VectorBundleFibers, p, X)
    Y = allocate_result(B, project_vector, p, X)
    return project_vector!(B, Y, p, X)
end

function project_vector!(B::VectorBundleFibers{<:TangentSpaceType}, Y, p, X)
    return project!(B.manifold, Y, p, X)
end
function project_vector!(B::VectorBundleFibers, Y, p, X)
    error("project_vector! not implemented for vector space family of type $(typeof(B)), output vector of type $(typeof(Y)) and input vector at point $(typeof(p)) with type of w $(typeof(X)).")
end

Base.@propagate_inbounds setindex!(x::FVector, val, i) = setindex!(x.data, val, i)

function representation_size(B::VectorBundleFibers{<:TCoTSpaceType})
    return representation_size(B.manifold)
end
function representation_size(B::VectorBundle)
    len_manifold = prod(representation_size(B.manifold))
    len_vs = prod(representation_size(B.fiber))
    return (len_manifold + len_vs,)
end

@doc raw"""
    sharp(M::Manifold, p, ξ::FVector)

Compute the sharp isomorphism (one of the musical isomorphisms) of vector `ξ`
from the vector space `M` at point `p` from the underlying [`Manifold`](@ref).

The function can be used for example to transform vectors
from the cotangent bundle to vectors from the tangent bundle
$♯ : T^{*}\mathcal M → T\mathcal M$
"""
function sharp(M::Manifold, p, ξ::FVector)
    X = allocate_result(M, sharp, ξ, p)
    return sharp!(M, X, p, ξ)
end

function sharp!(M::Manifold, X::FVector, p, ξ::FVector)
    error(
        "sharp! not implemented for vector bundle fibers space " *
        "of type $(typeof(M)), vector of type $(typeof(X)), point of " *
        "type $(typeof(p)) and vector of type $(typeof(ξ)).",
    )
end

@decorator_transparent_signature sharp!(
    M::AbstractDecoratorManifold,
    X::TFVector,
    p,
    ξ::CoTFVector,
)

show(io::IO, ::TangentSpaceType) = print(io, "TangentSpace")
show(io::IO, ::CotangentSpaceType) = print(io, "CotangentSpace")
function show(io::IO, tpt::TensorProductType)
    print(io, "TensorProductType(", join(tpt.spaces, ", "), ")")
end
function show(io::IO, fiber::VectorBundleFibers)
    print(io, "VectorBundleFibers($(fiber.fiber), $(fiber.manifold))")
end
function show(io::IO, mime::MIME"text/plain", vs::VectorSpaceAtPoint)
    summary(io, vs)
    println(io, "\nFiber:")
    pre = " "
    sf = sprint(show, "text/plain", vs.fiber; context = io, sizehint = 0)
    sf = replace(sf, '\n' => "\n$(pre)")
    println(io, pre, sf)
    println(io, "Base point:")
    sp = sprint(show, "text/plain", vs.point; context = io, sizehint = 0)
    sp = replace(sp, '\n' => "\n$(pre)")
    print(io, pre, sp)
end
show(io::IO, vb::VectorBundle) = print(io, "VectorBundle($(vb.type), $(vb.manifold))")
show(io::IO, vb::TangentBundle) = print(io, "TangentBundle($(vb.manifold))")
show(io::IO, vb::CotangentBundle) = print(io, "CotangentBundle($(vb.manifold))")

allocate(x::FVector) = FVector(x.type, allocate(x.data))
allocate(x::FVector, ::Type{T}) where {T} = FVector(x.type, allocate(x.data, T))

"""
    allocate_result(B::VectorBundleFibers, f, x...)

Allocates an array for the result of function `f` that is
an element of the vector space of type `B.fiber` on manifold `B.manifold`
and arguments `x...` for implementing the non-modifying operation
using the modifying operation.
"""
function allocate_result(B::VectorBundleFibers, f, x...)
    T = allocate_result_type(B, f, x)
    return allocate(x[1], T)
end
function allocate_result(M::Manifold, ::typeof(flat), w::TFVector, x)
    return FVector(CotangentSpace, allocate(w.data))
end
function allocate_result(M::Manifold, ::typeof(sharp), w::CoTFVector, x)
    return FVector(TangentSpace, allocate(w.data))
end

"""
    allocate_result_type(B::VectorBundleFibers, f, args::NTuple{N,Any}) where N

Returns type of element of the array that will represent the result of
function `f` for representing an operation with result in the vector space `fiber`
for manifold `M` on given arguments (passed at a tuple).
"""
function allocate_result_type(B::VectorBundleFibers, f, args::NTuple{N,Any}) where {N}
    T = typeof(reduce(+, one(number_eltype(eti)) for eti ∈ args))
    return T
end

size(x::FVector) = size(x.data)

function submanifold_component(M::Manifold, x::FVector, i::Val)
    return submanifold_component(M, x.data, i)
end
submanifold_component(x::FVector, i::Val) = submanifold_component(x.data, i)

submanifold_components(M::Manifold, x::FVector) = submanifold_components(M, x.data)
submanifold_components(x::FVector) = submanifold_components(x.data)

"""
    vector_space_dimension(B::VectorBundleFibers)

Dimension of the vector space of type `B`.
"""
function vector_space_dimension(B::VectorBundleFibers)
    error("vector_space_dimension not implemented for vector space family $(typeof(B)).")
end
function vector_space_dimension(B::VectorBundleFibers{<:TCoTSpaceType})
    return manifold_dimension(B.manifold)
end
function vector_space_dimension(B::VectorBundleFibers{<:TensorProductType})
    dim = 1
    for space in B.fiber.spaces
        dim *= vector_space_dimension(VectorBundleFibers(space, B.manifold))
    end
    return dim
end

"""
    zero_vector(B::VectorBundleFibers, p)

Compute the zero vector from the vector space of type `B.fiber` at point `p`
from manifold `B.manifold`.
"""
function zero_vector(B::VectorBundleFibers, p)
    X = allocate_result(B, zero_vector, p)
    return zero_vector!(B, X, p)
end

"""
    zero_vector!(B::VectorBundleFibers, X, p)

Save the zero vector from the vector space of type `B.fiber` at point `p`
from manifold `B.manifold` to `X`.
"""
function zero_vector!(B::VectorBundleFibers, X, p)
    error("zero_vector! not implemented for vector space family of type $(typeof(B)).")
end
function zero_vector!(B::VectorBundleFibers{<:TangentSpaceType}, X, p)
    return zero_tangent_vector!(B.manifold, X, p)
end

@doc raw"""
    zero_tangent_vector(B::VectorBundle, p)

Zero tangent vector at point `p` from the vector bundle `B`
over manifold `B.fiber` (denoted $\mathcal M$). The zero vector belongs to the space $T_{p}B$

Notation:
  * The point $p = (x_p, V_p)$ where $x_p ∈ \mathcal M$ and $V_p$ belongs to the
    fiber $F=π^{-1}(\{x_p\})$ of the vector bundle $B$ where $π$ is the
    canonical projection of that vector bundle $B$.

The zero vector is calculated as

$\mathbf{0}_{p} = (\mathbf{0}_{x_p}, \mathbf{0}_F)$

where $\mathbf{0}_{x_p}$ is the zero tangent vector from $T_{x_p}\mathcal M$ and
$\mathbf{0}_F$ is the zero element of the vector space $F$.
"""
zero_tangent_vector(::VectorBundle, ::Any...)

function zero_tangent_vector!(B::VectorBundle, X, p)
    xp, Vp = submanifold_components(B.manifold, p)
    VXM, VXF = submanifold_components(B.manifold, X)
    zero_tangent_vector!(B.manifold, VXM, xp)
    zero_vector!(B.fiber, VXF, Vp)
    return X
end


import LinearAlgebra: transpose
import Splittings: Geometry
import Splittings: SplineNN, interpolate!

@testset "Natural Spline" begin

    geom = Geometry( 101, 101, -5.0, -5.0, 0.1, 0.1 )
    x1, x2 = geom.x1grid, geom.x2grid
    dt = 0.5
    fc = exp.(-x1.^2) .* transpose(exp.(-x2.^2))
    spline = SplinePP( geom )
    interpolate!( spline, fc, dt, dt )
    fe = exp.(-(x1.-dt).^2) .* transpose(exp.(-(x2.-dt).^2))
    error = maximum( abs.( fe .- fc ))
    println( " error  = $error ")
    @test fc ≈ fe

    geom = Geometry( 51, 101, -5.0, -5.0, 0.2, 0.1 )
    x1, x2 = geom.x1grid, geom.x2grid
    dt = 0.5
    fc = exp.(-x1.^2) .* transpose(exp.(-x2.^2))
    spline = SplineNN( geom )
    interpolate!( spline, fc, dt, dt )
    fe = exp.(-(x1.-dt).^2) .* transpose(exp.(-(x2.-dt).^2))
    error = maximum( abs.( fe .- fc ))
    println( " error  = $error ")
    @test error < 1e-4

end


function transporter_density_to_monomer_uptake_sites(V_c::Vector{Float64}, ρ_closure::Matrix{Float64}, Min_gen_time::Vector{Float64}, Gram_stain::Vector{String})
    gmax = log(2)./Min_gen_time
    n_0 = cell_volume_to_cell_number_density(V_c, gmax, Gram_stain)
    r_p = 1e-9                    # 1 μm
    r_c = (3*V_c/(4*pi)).^(1/3)
    N_SB = zeros(size(ρ_closure,1), size(ρ_closure,2))
    for i in 1:size(V_c,1)
        N_SB[:,i] = n_0[i]*4*r_c[i]^2/(r_p^2)*ρ_closure[:,i]/12.011
    end
    N_SB[N_SB.<1e-12].=0.0
    return N_SB
end

function specific_reference_affinity(V_c::Vector{Float64}, ρ_p::Matrix{Float64}, D_S::Vector{Float64})
    # changed N_p (the number of transporters) to ρ_p = N_p*4*pi*r_p^2/(4*pi*r_c^2) (the transporter density on the cell surface)
    V_S = 1800.0                   # k2p, [s]
    N_A = 6.022e23                # Avogrado constant [1/mol]
    r_p = 1e-9                    # porter radius, 1 μm
    r_c = (3*V_c/(4*pi)).^(1/3)   # cell radius
    #
    p_inv = zeros(size(ρ_p,1), size(ρ_p,2))
    for i in 1:size(ρ_p,2)
        p_inv[:,i] = @. (4*ρ_p[:,i]*r_c[i]^2/r_p + pi*r_c[i])/(4*ρ_p[:,i]*r_c[i]^2/r_p)
    end
    #
    K_SC_0 = zeros(size(ρ_p,1), size(ρ_p,2))
    for i in 1:size(ρ_p,2)
        K_SC_0[:,i] = @. ((V_S*ρ_p[:,i]*r_c[i])/(pi*D_S*r_p^2*N_A))*p_inv[:,i]
    end
    #
    any(x->x==true, isnan.(K_SC_0)) ? throw(DomainError("NaN in DEBmicroTrait.ECA_kinetics!"))  : return K_SC_0
end

function transporters_closure(ρ_p::Vector{Float64}, genome_distr::Matrix{Float64})
    closure   = zeros(size(genome_distr))
    for i in 1:size(closure,2)
        closure[:,i] = ρ_p[i]*genome_distr[:,i]./sum(genome_distr[:,i])
    end
    closure[closure.==0.0].=1e-8
    return closure
end


# generate sample timeseries data
# a port of https://github.com/LeeDoYup/RobustSTL/blob/master/sample_generator.py
using StatsBase, Distributions

function getRandomChoice(total_length::Int, num_choice::Int;
                         replace = false)
    @assert !(replace == false && num_choice > total_length) "Draws without replacement must not exceed population size"
    sample(1:total_length, num_choice; replace = replace)
end

function getChange(change_type::String, change_level::Number)
    @assert change_type in ["random", "spike", "increase", "dip", "decrease"] "Not a valid change type"

    if change_type in ["random"]
        rand([-1, 1]) * change_level
    elseif change_type in ["dip", "decrease"]
        -1.0 * change_level
    else
        change_level
    end
end

function getSeason(season_n::Int, season_type::String, season_level::Number)
    @assert season_type in ["random", "stair"] "Not a valid season type"

    season = zeros(season_n)
    if season_type == "random"
        season = 2.0 * (rand(season_n) .- 0.5) * season_level
    else
        centre = season_n ÷ 2
        season[1:centre] .+= season_level
        season[centre + 1: end] .-= season_level
    end
    season
end

function genSeasons(total_n, season_n, season_num, season_type, season_level)
    season = getSeason(season_n, season_type, season_level)
    repeat(season, season_num)[1:total_n]
end

function genRemainder(total_n, noise_mean, noise_std)
    rand(Normal(noise_mean, noise_std), total_n)
end

function genAnomalies(total_n, anomaly_num, anomaly_type, anomaly_level)
    anomalyLocs = getRandomChoice(total_n, anomaly_num)
    anomalies = zeros(total_n)
    map(x -> anomalies[x] += getChange(anomaly_type, anomaly_level), anomalyLocs)
    anomalies
end

function genTrends(total_n, trend_change_num, trend_type, trend_level)
    trends = zeros(total_n)
    change_points = getRandomChoice(total_n, trend_change_num)
    map(x -> trends[x:end] .+= getChange(trend_type, trend_level), change_points)
    trends
end

"""
    gen_ts(total_n, season_n)

Generates a basic timeseries with trend, seasonality, noise and some anomalies.

**Parameters**:

- _total_n_: size of required dataset (e.g. 365, 52, 12, etc.).
- _season_n: number of seasonal points (e.g. 12 in case of a simulated monthly seasonality).

The following are optional parameters that can be passed to customise the individual timeseries components:

- _season_type_: set trajectory of seasonal component. `random` or `stair`. Default is `stair`.
- _season_level_: set amplitude of the seasonal effect (essentially acts as a static multiplier). Default is `1`.

- _trend_type_: set trajectory of trend component. `random` (default), `increase`, or `decrease`.
- _trend_level_: set amplitude of the trend component. Default is `3`
- _trend_change_num_: set how many instances of changes in the trend to be present. These are randomly distributed across the dataset. Default is `10`

- _anomaly_num_: set how many anomalies are to be randomly distributed across the dataset. Default is `6`.
- _anomaly_type_: set trajectory of the anomalies. `random` (default), `increase`, or `decrease`.
- _anomaly_level_: set amplitude of the anomalies. Default is `4`.

- _noise_mean_: mean of the normaly distributed noise component. Default is `0`.
- _noise_std_: standard deviation of the noise component. Default is `0.316`.

Returns a dataframe with individual components as well as the composite `sample`.
"""
function gen_ts(total_n::Number, season_n::Number;
    season_type = "stair",
    season_level = 1,
    trend_type = "random",
    trend_level = 3,
    trend_change_num = 10,
    anomaly_num = 6,
    anomaly_type = "random",
    anomaly_level = 4,
    noise_mean = 0,
    noise_std = 0.316)

    @assert total_n >= season_n "Dataset length must be equal or greater to seasonality"

    season_num = (total_n ÷ season_n) + 1

    trends = genTrends(total_n, trend_change_num, trend_type, trend_level)
    seasons = genSeasons(total_n, season_n, season_num, season_type, season_level)
    remainders = genRemainder(total_n, noise_mean, noise_std)
    anomalies = genAnomalies(total_n, anomaly_num, anomaly_type, anomaly_level)
    sample = trends + seasons + remainders + anomalies

    DataFrame(sample = sample,
            trends = trends,
            seasonality = seasons,
            remainder = remainders + anomalies)
end

"""
    gen_ts()

Generates a basic timeseries with trend, seasonality, noise and some anomalies.

Returns a dataframe with individual components as well as the composite `sample`.
"""
function gen_ts()
    gen_ts(365, 30)
end


using StatsBase
using Base.Test

X = randn(3, 8)

Z1 = X .- mean(X, 1)
Z2 = X .- mean(X, 2)

w1 = rand(3)
w2 = rand(8)

wv1 = weights(w1)
wv2 = weights(w2)

Z1w = X .- mean(X, wv1, 1)
Z2w = X .- mean(X, wv2, 2)

## reference results

S1 = Z1'Z1
S2 = Z2 * Z2'

Sz1 = X'X
Sz2 = X * X'

S1w = Z1w' * diagm(w1) * Z1w
S2w = Z2w * diagm(w2) * Z2w'

Sz1w = X' * diagm(w1) * X
Sz2w = X * diagm(w2) * X'

## scattermat

if VERSION < v"0.5.0-dev+679"
    @test scattermat(X)           ≈ S1
    @test scattermat(X; vardim=2) ≈ S2

    @test scattermat(X; mean=0)           ≈ Sz1
    @test scattermat(X; mean=0, vardim=2) ≈ Sz2

    @test scattermat(X; mean=mean(X,1))           ≈ S1
    @test scattermat(X; mean=mean(X,2), vardim=2) ≈ S2

    @test scattermat(X; mean=zeros(1,8))         ≈ Sz1
    @test scattermat(X; mean=zeros(3), vardim=2) ≈ Sz2

    ## weighted scatter mat

    @test scattermat(X, wv1)           ≈ S1w
    @test scattermat(X, wv2; vardim=2) ≈ S2w

    @test scattermat(X, wv1; mean=0)           ≈ Sz1w
    @test scattermat(X, wv2; mean=0, vardim=2) ≈ Sz2w

    @test scattermat(X, wv1; mean=mean(X, wv1, 1))           ≈ S1w
    @test scattermat(X, wv2; mean=mean(X, wv2, 2), vardim=2) ≈ S2w

    @test scattermat(X, wv1; mean=zeros(1,8))         ≈ Sz1w
    @test scattermat(X, wv2; mean=zeros(3), vardim=2) ≈ Sz2w
else
    @test scattermat(X)    ≈ S1
    @test scattermat(X, 2) ≈ S2

    @test StatsBase.scattermatm(X, 0)    ≈ Sz1
    @test StatsBase.scattermatm(X, 0, 2) ≈ Sz2

    @test StatsBase.scattermatm(X, mean(X,1))    ≈ S1
    @test StatsBase.scattermatm(X, mean(X,2), 2) ≈ S2

    @test StatsBase.scattermatm(X, zeros(1,8))  ≈ Sz1
    @test StatsBase.scattermatm(X, zeros(3), 2) ≈ Sz2

    ## weighted scatter mat

    @test scattermat(X, wv1)    ≈ S1w
    @test scattermat(X, wv2, 2) ≈ S2w

    @test StatsBase.scattermatm(X, 0, wv1)    ≈ Sz1w
    @test StatsBase.scattermatm(X, 0, wv2, 2) ≈ Sz2w

    @test StatsBase.scattermatm(X, mean(X, wv1, 1), wv1)    ≈ S1w
    @test StatsBase.scattermatm(X, mean(X, wv2, 2), wv2, 2) ≈ S2w

    @test StatsBase.scattermatm(X, zeros(1,8), wv1)  ≈ Sz1w
    @test StatsBase.scattermatm(X, zeros(3), wv2, 2) ≈ Sz2w
end

# weighted covariance

if VERSION < v"0.5.0-dev+679"
    @test cov(X, wv1)           ≈ S1w ./ sum(wv1)
    @test cov(X, wv2; vardim=2) ≈ S2w ./ sum(wv2)

    @test cov(X, wv1; mean=0)           ≈ Sz1w ./ sum(wv1)
    @test cov(X, wv2; mean=0, vardim=2) ≈ Sz2w ./ sum(wv2)

    @test cov(X, wv1; mean=mean(X, wv1, 1))           ≈ S1w ./ sum(wv1)
    @test cov(X, wv2; mean=mean(X, wv2, 2), vardim=2) ≈ S2w ./ sum(wv2)

    @test cov(X, wv1; mean=zeros(1,8))         ≈ Sz1w ./ sum(wv1)
    @test cov(X, wv2; mean=zeros(3), vardim=2) ≈ Sz2w ./ sum(wv2)
else
    @test cov(X, wv1)    ≈ S1w ./ sum(wv1)
    @test cov(X, wv2, 2) ≈ S2w ./ sum(wv2)

    @test Base.covm(X, 0, wv1)    ≈ Sz1w ./ sum(wv1)
    @test Base.covm(X, 0, wv2, 2) ≈ Sz2w ./ sum(wv2)

    @test Base.covm(X, mean(X, wv1, 1), wv1)    ≈ S1w ./ sum(wv1)
    @test Base.covm(X, mean(X, wv2, 2), wv2, 2) ≈ S2w ./ sum(wv2)

    @test Base.covm(X, zeros(1,8), wv1)  ≈ Sz1w ./ sum(wv1)
    @test Base.covm(X, zeros(3), wv2, 2) ≈ Sz2w ./ sum(wv2)
end

# mean_and_cov
if VERSION < v"0.5.0-dev+679"
    (m, C) = mean_and_cov(X; vardim=1)
    @test m == mean(X, 1)
    @test C == cov(X; vardim=1)

    (m, C) = mean_and_cov(X; vardim=2)
    @test m == mean(X, 2)
    @test C == cov(X; vardim=2)

    (m, C) = mean_and_cov(X, wv1; vardim=1)
    @test m == mean(X, wv1, 1)
    @test C == cov(X, wv1; vardim=1)

    (m, C) = mean_and_cov(X, wv2; vardim=2)
    @test m == mean(X, wv2, 2)
    @test C == cov(X, wv2; vardim=2)
else
    (m, C) = mean_and_cov(X, 1)
    @test m == mean(X, 1)
    @test C == cov(X, 1)

    (m, C) = mean_and_cov(X, 2)
    @test m == mean(X, 2)
    @test C == cov(X, 2)

    (m, C) = mean_and_cov(X, wv1, 1)
    @test m == mean(X, wv1, 1)
    @test C == cov(X, wv1, 1)

    (m, C) = mean_and_cov(X, wv2, 2)
    @test m == mean(X, wv2, 2)
    @test C == cov(X, wv2, 2)
end


using Test

using NeXLCore

@testset "Stoichiometry" begin

    # Check against DTSA-II
    @test isapprox(
        obystoichiometry(
            n"Mg" => 0.116567,
            n"Al" => 0.0490615,
            n"Si" => 0.211982,
            n"Ca" => 0.10899,
            n"Fe" => 0.0774196,
        ),
        0.4276,
        atol = 0.0001,
    )

    sio2 = asoxide(n"Si")
    @test name(sio2) == "SiO₂"
    asio2 = atomicfraction("SiO2", n"Si" => 1, n"O" => 2)
    asio2 = atomicfraction("SiO₂", n"Si" => 1, n"O" => 2)
    @test isapprox(sio2[n"Si"], asio2[n"Si"], atol = 1.0e-6)
    @test isapprox(sio2[n"O"], asio2[n"O"], atol = 1.0e-6)

    mix = sum(asoxide(
        n"Mg" => 0.116567,
        n"Al" => 0.0490615,
        n"Si" => 0.211982,
        n"Ca" => 0.10899,
        n"Fe" => 0.0774196))

    @test isapprox(mix[n"O"], 0.4276, atol = 0.0001)
    @test isapprox(mix[n"Al"], 0.0490615, atol = 0.0001)
    @test isapprox(mix[n"Fe"], 0.0774196, atol = 0.0001)

    ox = sum(asoxide(massfraction(delete(srm470_k412, n"O"))), name="K412")
    @test all(isapprox(ox[elm], srm470_k412[elm], atol=1.0e-8) for elm in keys(srm470_k412))

    ox = sum(asoxide(massfraction(delete(srm470_k411, [ n"O" ]))), name="K411")
    @test all(isapprox(ox[elm], srm470_k411[elm], atol=1.0e-8) for elm in keys(srm470_k411))

    m=sum(asoxide(n"Al"=>0.1029, n"Na"=>0.0877, n"Si"=>0.3213), name="Albite", density=5.0, description="a common mineral", pedigree="Natural")
    @test isapprox(m[n"O"], 0.4881, atol=0.0001)
    @test isapprox(m[n"Si"], 0.3212, atol=0.0001)
    @test isapprox(m[n"Na"], 0.0877, atol=0.0001)
    @test m[:Density] == 5.0
    @test m[:Description] == "a common mineral"
    @test m[:Pedigree] == "Natural"
    @test name(m) == "Albite"

    amix =
        0.4 * atomicfraction("SiO2", n"Si" => 1, n"O" => 2) +
        0.6 * atomicfraction("Al2O3", n"Al" => 2, n"O" => 3)
    @test isapprox(amix[n"O"], 0.4955, atol = 0.0001)
    @test isapprox(amix[n"Al"], 0.3176, atol = 0.0001)
    @test isapprox(amix[n"Si"], 0.1870, atol = 0.0001)

    # Check against DTSA-II
    @test isapprox(amix[n"Si"], 0.1870, atol = 0.0001)
    @test isapprox(amix[n"Al"], 0.3176, atol = 0.0001)
    @test isapprox(amix[n"O"], 0.4955, atol = 0.0001)
end


using QOptCtrl

X = pauli_X
Y = pauli_Y
Z = pauli_Z
H = 1/sqrt(2)*[1 1; 1 -1] # Hadamard Gate

ω = 1

prob = GateTransfer(
        (ω/2)*Z, # drift Hamiltonian
        [X, Y], # control Hamiltonians
        20, # number of steps
        4, # transfer time
        H, # tarfet gate
        rand(20, 2) #initial guess
)

init_ctrl_arr,  final_ctrl_arr = OPTIMIZE(prob)
visualise_pulses(init_ctrl_arr, final_ctrl_arr)


struct ClusterWolff end
"""
Wolff single cluster update

#Arguments:
* `spins` : array of spin values
* `nearest_neighbors` : function that returns a list of nearest neighbors to index i
* `beta` : inverse temperature
* `rng`  : random number generator

Ref: U. Wolff, Phys. Rev. Lett. 62, 361 (1989)

consider for Potts:
https://stanford.edu/~kurinsky/ClassWork/Physics271_Final_Paper.pdf

consider BlumeCapel implementation

todo:
return dE?
how to specify function type for performance?
"""
function update(alg::ClusterWolff, rng::AbstractRNG, spins::Vector{Int}, nearest_neighbors::Function, betaJ::Float64)
    N = length(spins)

    # get random spin and flip it
    index_i = rand(rng, 1:N)
    s_i = spins[index_i]
    spins[index_i] *= -1

    # probability (per bond!)
    p = 1.0 - exp(-2.0 * betaJ)

    # check all bonds and already flip cluster
    proposed = copy(nearest_neighbors(index_i))
    while !isempty(proposed)
        # check proposed sites for alignment with base spin s_i and add to cluster with probability p
        index_j = pop!(proposed)
        # no need to check if site is already in cluster, because spin flip disqualifies it already
        if spins[index_j] == s_i && rand(rng) < p
            spins[index_j] *= -1
            append!(proposed, nearest_neighbors(index_j))
        end
    end
end


# VSMBasedUnwarping.jl
# Methods for unwarping geometric distortions in MRI images

function unwarp(VSM, distorted, dim)
    if dim == 2
        distorted = switchdim(distorted)
        VSM = switchdim(VSM)
    end

    unwarped = unwarp(VSM, distorted)

    if dim == 2
        unwarped = switchdim(unwarped)
    end
    unwarped
end

switchdim(v) = permutedims(v, [2, 1, (3:ndims(v))...])

unwarp(VSM, distorted) = unwarp!(similar(distorted), VSM, distorted)

function unwarp!(unwarped, VSM, distorted)
    xi = 1:size(distorted, 1)
    for J in CartesianIndices(size(distorted)[4:end])
        for I in CartesianIndices(size(distorted)[2:3])
            xtrue = xi .+ VSM[:,I]
            xregrid = (xtrue[1] .<= xi .<= xtrue[end]) # only use x values inside (no extrapolation)
            unwarped[.!xregrid,I,J] .= 0
            unwarped[xregrid,I,J] .= unwarpline(xtrue, distorted[:,I,J], xi[xregrid])
        end
    end
    unwarped
end

function unwarpline(xtrue, distorted, xnew)
    #TODO try better interpolation than linear
    interpolate((xtrue,), distorted, Gridded(Linear()))(xnew)
end

function getVSM(B0, rbw, dim, threshold = 5.0)
    VSM = B0 ./ (2π * rbw)
    thresholdforward(VSM, -0.9, threshold, dim)
end

function thresholdforward(VSM, tmin, tmax, dim)
    if dim == 2
        VSM = switchdim(VSM)
    end

    deltaVSM = VSM[2:end,:,:] .- VSM[1:(end-1),:,:]
    VSMret = copy(VSM)
    nx = size(VSM, 1)
    for I in CartesianIndices(size(VSM)[2:3])
        for x in 1:(nx-1)
            if deltaVSM[x,I] < tmin || deltaVSM[x,I] > tmax
                if deltaVSM[x,I] < tmin
                    diff = tmin - deltaVSM[x,I]
                else
                    diff = tmax - deltaVSM[x,I]
                end
                VSMret[x+1,I] += diff
                if x + 1 < nx
                    deltaVSM[x+1,I] -= diff
                end
            end
        end
    end

    if dim == 2
        VSMret = switchdim(VSMret)
    end
    VSMret
end

# TODO not properly working
threshold(VSM, tmin, tmax, dim) = (thresholdforward(VSM, tmin, tmax, dim) .+ thresholdforward(VSM[end:-1:1,:,:], -tmax, -tmin, dim))[end:-1:1,:,:] ./ 2.0