# Naive implementations of multiply_dimensionwise used to demonstrate the functionality
# without performance optimizations and for testing correctness of the optimized versions
# implemented below.
function multiply_dimensionwise_naive(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 2})
  size_out = size(matrix, 1)
  size_in  = size(matrix, 2)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out)

  for i in 1:size_out
    for ii in 1:size_in
      for v in 1:n_vars
        data_out[v, i] += matrix[i, ii] * data_in[v, ii]
      end
    end
  end

  return data_out
end

function multiply_dimensionwise_naive(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 3})
  size_out = size(matrix, 1)
  size_in  = size(matrix, 2)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out, size_out)

  for j in 1:size_out, i in 1:size_out
    for jj in 1:size_in, ii in 1:size_in
      for v in 1:n_vars
        data_out[v, i, j] += matrix[i, ii] * matrix[j, jj] * data_in[v, ii, jj]
      end
    end
  end

  return data_out
end

function multiply_dimensionwise_naive(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 4})
  size_out = size(matrix, 1)
  size_in  = size(matrix, 2)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out, size_out, size_out)

  for k in 1:size_out, j in 1:size_out, i in 1:size_out
    for kk in 1:size_in, jj in 1:size_in, ii in 1:size_in
      for v in 1:n_vars
        data_out[v, i, j, k] += matrix[i, ii] * matrix[j, jj] * matrix[k, kk] * data_in[v, ii, jj, kk]
      end
    end
  end

  return data_out
end

"""
    multiply_dimensionwise(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, NDIMS+1})

Multiply the array `data_in` by `matrix` in each coordinate direction, where `data_in`
is assumed to have the first coordinate for the number of variables and the remaining coordinates
are multiplied by `matrix`.
"""
function multiply_dimensionwise(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 2})
  # 1D
  # optimized version of multiply_dimensionwise_naive
  size_out = size(matrix, 1)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out)

  multiply_dimensionwise!(data_out, matrix, data_in)

  return data_out
end

function multiply_dimensionwise(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 3})
  # 2D
  # optimized version of multiply_dimensionwise_naive
  size_out = size(matrix, 1)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out, size_out)

  multiply_dimensionwise!(data_out, matrix, data_in)

  return data_out
end

function multiply_dimensionwise(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 4})
  # 3D
  # optimized version of multiply_dimensionwise_naive
  size_out = size(matrix, 1)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out, size_out, size_out)

  multiply_dimensionwise!(data_out, matrix, data_in)

  return data_out
end

# In the following, there are several optimized in-place versions of multiply_dimensionwise.
# These make use of the macro `@tullio` from Tullio.jl, which basically uses an Einstein
# summation convention syntax.

# 1D version
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 2}, matrix::AbstractMatrix,
                                 data_in ::AbstractArray{<:Any, 2})
  @tullio threads=false data_out[v, i] = matrix[i, ii] * data_in[v, ii]

  return nothing
end

# 1D version, apply matrixJ to data_inJ
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 2}, matrix1::AbstractMatrix,
                                 data_in1::AbstractArray{<:Any, 2}, matrix2::AbstractMatrix,
                                 data_in2::AbstractArray{<:Any, 2})
  @tullio threads=false data_out[v, i] = matrix1[i, ii] * data_in1[v, ii] + matrix2[i, ii] * data_in2[v, ii]

  return nothing
end

# 2D version
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 3}, matrix::AbstractMatrix,
                                 data_in:: AbstractArray{<:Any, 3},
                                 tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix, 1), size(matrix, 2)))

  # Interpolate in x-direction
  @tullio threads=false tmp1[v, i, j]     = matrix[i, ii] * data_in[v, ii, j]

  # Interpolate in y-direction
  @tullio threads=false data_out[v, i, j] = matrix[j, jj] * tmp1[v, i, jj]

  return nothing
end

# 2D version, apply matrixJ to dimension J of data_in
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 3},
                                 matrix1::AbstractMatrix, matrix2::AbstractMatrix,
                                 data_in:: AbstractArray{<:Any, 3},
                                 tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 2)))

  # Interpolate in x-direction
  @tullio threads=false tmp1[v, i, j]     = matrix1[i, ii] * data_in[v, ii, j]

  # Interpolate in y-direction
  @tullio threads=false data_out[v, i, j] = matrix2[j, jj] * tmp1[v, i, jj]

  return nothing
end

# 2D version, apply matrixJ to dimension J of data_in and add the result to data_out
function add_multiply_dimensionwise!(data_out::AbstractArray{<:Any, 3},
                                     matrix1::AbstractMatrix, matrix2::AbstractMatrix,
                                     data_in:: AbstractArray{<:Any, 3},
                                     tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 2)))

  # Interpolate in x-direction
  @tullio threads=false tmp1[v, i, j]     = matrix1[i, ii] * data_in[v, ii, j]

  # Interpolate in y-direction
  @tullio threads=false data_out[v, i, j] += matrix2[j, jj] * tmp1[v, i, jj]

  return nothing
end

# 3D version
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 4}, matrix::AbstractMatrix,
                                 data_in:: AbstractArray{<:Any, 4},
                                 tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix, 1), size(matrix, 2), size(matrix, 2)),
                                 tmp2=zeros(eltype(data_out), size(data_out, 1), size(matrix, 1), size(matrix, 1), size(matrix, 2)))

  # Interpolate in x-direction
  @tullio threads=false tmp1[v, i, j, k]     = matrix[i, ii] * data_in[v, ii, j, k]

  # Interpolate in y-direction
  @tullio threads=false tmp2[v, i, j, k]     = matrix[j, jj] * tmp1[v, i, jj, k]

  # Interpolate in z-direction
  @tullio threads=false data_out[v, i, j, k] = matrix[k, kk] * tmp2[v, i, j, kk]

  return nothing
end

# 3D version, apply matrixJ to dimension J of data_in
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 4},
                                 matrix1::AbstractMatrix, matrix2::AbstractMatrix, matrix3::AbstractMatrix,
                                 data_in:: AbstractArray{<:Any, 4},
                                 tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 2), size(matrix1, 2)),
                                 tmp2=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 1), size(matrix1, 2)))

  # Interpolate in x-direction
  @tullio threads=false tmp1[v, i, j, k]     = matrix1[i, ii] * data_in[v, ii, j, k]

  # Interpolate in y-direction
  @tullio threads=false tmp2[v, i, j, k]     = matrix2[j, jj] * tmp1[v, i, jj, k]

  # Interpolate in z-direction
  @tullio threads=false data_out[v, i, j, k] = matrix3[k, kk] * tmp2[v, i, j, kk]

  return nothing
end

# 3D version, apply matrixJ to dimension J of data_in and add the result to data_out
function add_multiply_dimensionwise!(data_out::AbstractArray{<:Any, 4},
                                     matrix1::AbstractMatrix, matrix2::AbstractMatrix, matrix3::AbstractMatrix,
                                     data_in:: AbstractArray{<:Any, 4},
                                     tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 2), size(matrix1, 2)),
                                     tmp2=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 1), size(matrix1, 2)))

  # Interpolate in x-direction
  @tullio threads=false tmp1[v, i, j, k]     = matrix1[i, ii] * data_in[v, ii, j, k]

  # Interpolate in y-direction
  @tullio threads=false tmp2[v, i, j, k]     = matrix2[j, jj] * tmp1[v, i, jj, k]

  # Interpolate in z-direction
  @tullio threads=false data_out[v, i, j, k] += matrix3[k, kk] * tmp2[v, i, j, kk]

  return nothing
end

# Calculate the Dhat matrix
function calc_dhat(nodes, weights)
  n_nodes = length(nodes)
  dhat = polynomial_derivative_matrix(nodes)
  dhat = transpose(dhat)

  for n in 1:n_nodes, j in 1:n_nodes
    dhat[j, n] *= -weights[n] / weights[j]
  end

  return dhat
end

# Calculate the Dsplit matrix for split-form differentiation: dplit = 2D - M⁻¹B
function calc_dsplit(nodes, weights)
  # Start with 2 x the normal D matrix
  dsplit = polynomial_derivative_matrix(nodes)
  dsplit = 2 .* dsplit

  # Modify to account for
  dsplit[1, 1] += 1/weights[1]
  dsplit[end, end] -= 1/weights[end]

  return dsplit
end

# Calculate the polynomial derivative matrix D
function polynomial_derivative_matrix(nodes)
  n_nodes = length(nodes)
  d = zeros(n_nodes, n_nodes)
  wbary = barycentric_weights(nodes)

  for i in 1:n_nodes, j in 1:n_nodes
    if j != i
      d[i, j] = wbary[j] / wbary[i] * 1 / (nodes[i] - nodes[j])
      d[i, i] -= d[i, j]
    end
  end

  return d
end

# Calculate and interpolation matrix (Vandermonde matrix) between two given sets of nodes
function polynomial_interpolation_matrix(nodes_in, nodes_out)
  n_nodes_in = length(nodes_in)
  n_nodes_out = length(nodes_out)
  wbary_in = barycentric_weights(nodes_in)
  vdm = zeros(n_nodes_out, n_nodes_in)

  for k in 1:n_nodes_out
    match = false
    for j in 1:n_nodes_in
      if isapprox(nodes_out[k], nodes_in[j], rtol=eps())
        match = true
        vdm[k, j] = 1
      end
    end

    if match == false
      s = 0.0
      for j in 1:n_nodes_in
        t = wbary_in[j] / (nodes_out[k] - nodes_in[j])
        vdm[k, j] = t
        s += t
      end
      for j in 1:n_nodes_in
        vdm[k, j] = vdm[k, j] / s
      end
    end
  end

  return vdm
end

# Calculate the barycentric weights for a given node distribution.
function barycentric_weights(nodes)
  n_nodes = length(nodes)
  weights = ones(n_nodes)

  for j = 2:n_nodes, k = 1:(j-1)
    weights[k] *= nodes[k] - nodes[j]
    weights[j] *= nodes[j] - nodes[k]
  end

  for j in 1:n_nodes
    weights[j] = 1 / weights[j]
  end

  return weights
end

# Calculate Lhat.
function calc_lhat(x::Float64, nodes, weights)
  n_nodes = length(nodes)
  wbary = barycentric_weights(nodes)

  lhat = lagrange_interpolating_polynomials(x, nodes, wbary)

  for i in 1:n_nodes
    lhat[i] /= weights[i]
  end

  return lhat
end

# Calculate Lagrange polynomials for a given node distribution.
function lagrange_interpolating_polynomials(x::Float64, nodes, wbary)
  n_nodes = length(nodes)
  polynomials = zeros(n_nodes)

  for i in 1:n_nodes
    if isapprox(x, nodes[i], rtol=eps(x))
      polynomials[i] = 1
      return polynomials
    end
  end

  for i in 1:n_nodes
    polynomials[i] = wbary[i] / (x - nodes[i])
  end
  total = sum(polynomials)

  for i in 1:n_nodes
    polynomials[i] /= total
  end

  return polynomials
end

# From FLUXO (but really from blue book by Kopriva)
function gauss_lobatto_nodes_weights(n_nodes::Integer)
  # From Kopriva's book
  n_iterations = 10
  tolerance = 1e-15

  # Initialize output
  nodes = zeros(n_nodes)
  weights = zeros(n_nodes)

  # Get polynomial degree for convenience
  N = n_nodes - 1

  # Calculate values at boundary
  nodes[1] = -1.0
  nodes[end] = 1.0
  weights[1] = 2 / (N * (N + 1))
  weights[end] = weights[1]

  # Calculate interior values
  if N > 1
    cont1 = pi/N
    cont2 = 3/(8 * N * pi)

    # Use symmetry -> only left side is computed
    for i in 1:(div(N + 1, 2) - 1)
      # Calculate node
      # Initial guess for Newton method
      nodes[i+1] = -cos(cont1*(i+0.25) - cont2/(i+0.25))

      # Newton iteration to find root of Legendre polynomial (= integration node)
      for k in 0:n_iterations
        q, qder, _ = calc_q_and_l(N, nodes[i+1])
        dx = -q/qder
        nodes[i+1] += dx
        if abs(dx) < tolerance * abs(nodes[i+1])
          break
        end
      end

      # Calculate weight
      _, _, L = calc_q_and_l(N, nodes[i+1])
      weights[i+1] = weights[1] / L^2

      # Set nodes and weights according to symmetry properties
      nodes[N+1-i] = -nodes[i+1]
      weights[N+1-i] = weights[i+1]
    end
  end

  # If odd number of nodes, set center node to origin (= 0.0) and calculate weight
  if n_nodes % 2 == 1
    _, _, L = calc_q_and_l(N, 0)
    nodes[div(N, 2) + 1] = 0.0
    weights[div(N, 2) + 1] = weights[1] / L^2
  end

  return nodes, weights
end

# From FLUXO (but really from blue book by Kopriva)
function calc_q_and_l(N::Integer, x::Float64)
  L_Nm2 = 1.0
  L_Nm1 = x
  Lder_Nm2 = 0.0
  Lder_Nm1 = 1.0

  local L
  for i in 2:N
    L = ((2 * i - 1) * x * L_Nm1 - (i - 1) * L_Nm2) / i
    Lder = Lder_Nm2 + (2 * i - 1) * L_Nm1
    L_Nm2 = L_Nm1
    L_Nm1 = L
    Lder_Nm2 = Lder_Nm1
    Lder_Nm1 = Lder
  end

  q = (2 * N + 1)/(N + 1) * (x * L - L_Nm2)
  qder = (2 * N + 1) * L

  return q, qder, L
end
calc_q_and_l(N::Integer, x::Real) = calc_q_and_l(N, convert(Float64, x))

# From FLUXO (but really from blue book by Kopriva)
function gauss_nodes_weights(n_nodes::Integer)
  # From Kopriva's book
  n_iterations = 10
  tolerance = 1e-15

  # Initialize output
  nodes = ones(n_nodes) * 1000
  weights = zeros(n_nodes)

  # Get polynomial degree for convenience
  N = n_nodes - 1
  if N == 0
    nodes .= 0.0
    weights .= 2.0
    return nodes, weights
  elseif N == 1
    nodes[1] = -sqrt(1/3)
    nodes[end] = -nodes[1]
    weights .= 1.0
    return nodes, weights
  else # N > 1
    # Use symmetry property of the roots of the Legendre polynomials
    for i in 0:(div(N + 1, 2) - 1)
      # Starting guess for Newton method
      nodes[i+1] = -cos(pi / (2 * N + 2) * (2 * i + 1))

      # Newton iteration to find root of Legendre polynomial (= integration node)
      for k in 0:n_iterations
        poly, deriv = legendre_polynomial_and_derivative(N + 1, nodes[i+1])
        dx = -poly / deriv
        nodes[i+1] += dx
        if abs(dx) < tolerance * abs(nodes[i+1])
          break
        end
      end

      # Calculate weight
      poly, deriv = legendre_polynomial_and_derivative(N + 1, nodes[i+1])
      weights[i+1] = (2 * N + 3) / ((1 - nodes[i+1]^2) * deriv^2)

      # Set nodes and weights according to symmetry properties
      nodes[N+1-i] = -nodes[i+1]
      weights[N+1-i] = weights[i+1]
    end

    # If odd number of nodes, set center node to origin (= 0.0) and calculate weight
    if n_nodes % 2 == 1
      poly, deriv = legendre_polynomial_and_derivative(N + 1, 0.0)
      nodes[div(N, 2) + 1] = 0.0
      weights[div(N, 2) + 1] = (2 * N + 3) / deriv^2
    end

    return nodes, weights
  end
end

# From FLUXO (but really from blue book by Kopriva)
function legendre_polynomial_and_derivative(N::Int, x::Real)
  if N == 0
    poly = 1.0
    deriv = 0.0
  elseif N == 1
    poly = convert(Float64, x)
    deriv = 1.0
  else
    poly_Nm2 = 1.0
    poly_Nm1 = convert(Float64, x)
    deriv_Nm2 = 0.0
    deriv_Nm1 = 1.0

    poly = 0.0
    deriv = 0.0
    for i in 2:N
      poly = ((2*i-1) * x * poly_Nm1 - (i-1) * poly_Nm2) / i
      deriv=deriv_Nm2 + (2*i-1)*poly_Nm1
      poly_Nm2=poly_Nm1
      poly_Nm1=poly
      deriv_Nm2=deriv_Nm1
      deriv_Nm1=deriv
    end
  end

  # Normalize
  poly = poly * sqrt(N+0.5)
  deriv = deriv * sqrt(N+0.5)

  return poly, deriv
end

# Calculate Legendre vandermonde matrix and its inverse
function vandermonde_legendre(nodes, N)
  n_nodes = length(nodes)
  n_modes = N + 1
  vandermonde = zeros(n_nodes, n_modes)

  for i in 1:n_nodes
    for m in 1:n_modes
      vandermonde[i, m], _ = legendre_polynomial_and_derivative(m-1, nodes[i])
    end
  end
  # for very high polynomial degree, this is not well conditioned
  inverse_vandermonde = inv(vandermonde)
  return vandermonde, inverse_vandermonde
end
vandermonde_legendre(nodes) = vandermonde_legendre(nodes, length(nodes) - 1)


module Mazes
export Maze, generate_btree!, generate_sidewinder!

using Graphs, Random

struct Maze
    dim::Tuple # TODO
    graph::SimpleGraph
end
function Maze(dim)
    graph = SimpleGraph(prod(dim))
    return Maze(dim, graph)
end
function Base.show(io::IO, maze::Maze)
    if length(maze.dim) == 2
        empty = "   "
        rows, cols = maze.dim
        println(io, "+" * "---+"^cols)
        for row = 1:rows
            sides = "|"
            floors = "+"
            for col = 1:cols
                # julia uses column major order
                cell = cols * (col - 1) + row
                right = cols * (col) + row
                down = cols * (col - 1) + row + 1
                sides = sides * empty * (
                    if has_edge(maze.graph, cell, right)
                        " "
                    else
                        "|"
                    end
                )
                floors = floors * (
                    if has_edge(maze.graph, cell, down)
                        "   "
                    else
                        "---"
                    end
                ) * "+"
            end
            println(io, sides)
            println(io, floors)
        end
    else
        print("Higher dimensions are not implemented.")
    end
end

function generate_btree!(maze)
    if length(maze.dim) == 2
        rows, cols = maze.dim
        for cell in 1:prod(maze.dim)
            row = mod(cell-1, rows) + 1
            col = fld(cell-1, rows) + 1
            if row + 1 > rows
                linked = cols * (col) + row
            elseif col + 1 > cols
                linked = cols * (col - 1) + row + 1
            else
                down, right = shuffle([0,1])
                linked = cols * (col + right - 1) + row + down
            end
            add_edge!(maze.graph, cell, linked)
        end
    else
        print("Higher dimensions are not implemented.")
    end
end

function generate_sidewinder!(maze) # BUG
    if length(maze.dim) == 2
        rows, cols = maze.dim
        run = []
        for cell in 1:prod(maze.dim)
            row = mod(cell-1, rows) + 1
            col = fld(cell-1, rows) + 1
            if row + 1 > rows
                linked = cols * (col) + row
            elseif col + 1 > cols
                linked = cols * (col - 1) + row + 1
            else
                append!(run, cell)
                forward = rand([true, false])
                if forward
                    linked = cols * (col - 1) + row + 1
                else
                    choice = rand(run)
                    add_edge!(maze.graph, choice, choice+1)
                    empty!(run)
                    continue
                end
            end
            add_edge!(maze.graph, cell, linked)
        end
    else
        print("Higher dimensions are note implemented.")
    end
end

end


Base.rand(m::SequentialPointProcess, window, npoints::Integer; giveup=-1, returnanyway=false) =
  rpattern(m, window, npoints, giveup, returnanyway)
function rpattern(model, window, npoints, giveup=-1, returnanyway=false)
  xy = Vector{NTuple{2,Float64}}(undef, npoints)
  if npoints == 0
    return PointPattern(xy, window)
  end
  xy[1] = randpoint(Uniform(), xy, window, giveup)
  for i in 2:npoints
    prop = randpoint(model, view(xy, 1:i-1), window, giveup)
    if isnothing(prop) # If sampling from the sequential model fails at some point return nothing.
      # This works with ABC, since that kind of parameter configuration is too unlikely.
      # In other contexts, giveup is not used.
      if returnanyway
        return PointPattern(xy[1:i], window)
      else
        return PointPattern(empty(xy), window)
      end
    end
    xy[i] = prop
  end
  PointPattern(xy, window)
end
function acceptpoint(m::Softcore, xyprop, xyold)
  p = fk(m, xyprop, xyold)
  return rand() <= p
end
function acceptpoint(m::OverlappingDiscs, xyprop, xyold)
  k = length(xyold) + 1
  p = m.theta(countoverlaps(m, xyprop, xyold), k) / maximum(n -> m.theta(n,k), 0:length(xyold))
  return rand() <= p
end
function acceptpoint(::Uniform, xyprop, xyold)
  true
end
function randpoint(m::Mixture, pattern, window, giveup)
  if rand() < m.theta
    randpoint(m.first, pattern, window, giveup)
  else
    randpoint(m.second, pattern, window, giveup)
  end
end
function randp(window)
  randcoord(x) = window[x][1] + (window[x][2]-window[x][1])*rand()
  randx() = randcoord(:x)
  randy() = randcoord(:y)
  randx(), randy()
end
function randpoint(model, pattern, window, giveup=-1)
  tries = 0
  while true
    xyprop = randp(window)
    if acceptpoint(model, xyprop, pattern)
      return xyprop
    end
    tries += 1
    if giveup != -1 && tries > giveup
      return nothing
    end
  end
  return nothing
end


using ReversibleSeismic, Test

@testset "loss" begin
     nx = 100
     ny = 100
     param = AcousticPropagatorParams(nx=nx, ny=ny,
          Rcoef=0.2, dx=20.0, dy=20.0, dt=0.05, nstep=2000)

     c = 1000*ones(param.NX+2, param.NY+2)
     src = (param.NX÷2, param.NY÷2)
     srcv = Ricker(param, 100.0, 500.0)
     tu = solve(param, src, srcv, c)
     tu2 = ReversibleSeismic.solve_final(param, src, srcv, c)
     tu3 = solve2(param, src, srcv, c)
     @test tu[:,:,end] ≈ tu2
     @test tu[:,:,end] ≈ tu3.u
     loss = sum(tu .^ 2)
     @test loss ≈ 10.931466822080788
end


import Base.-

-(A::Assoc,B::Assoc) = minus(A::Assoc,B::Assoc)
function minus(A::Assoc,B::Assoc)
    #Check A,B, if string => Logical
    At = A
    Bt = B
    if(A.val != [1.0])
        At = logical(A)
    end
    if(B.val !=[1.0])
        Bt = logical(B)
    end
    ## A*B operation
    ABrow = sortedunion(At.row,Bt.row)
    ABcol = sortedunion(At.col,Bt.col)
    if (size(ABrow,1) == 0 || size(ABcol,1) == 0)
        return Assoc([1],[1],0,(+))
    end

    ABA =  spzeros(length(ABrow),length(ABcol))
    Arow = searchsortedmapping(At.row,ABrow)
    Acol = searchsortedmapping(At.col,ABcol)
    Brow = searchsortedmapping(Bt.row,ABrow)
    Bcol = searchsortedmapping(Bt.col,ABcol)
    ABA[Arow,Acol] += At.A
    ABA[Brow,Bcol] -= Bt.A
    AB = Assoc(ABrow,ABcol, Array{Union{AbstractString,Number}}([1.0]), ABA)

    return condense(AB) #Incase of negation
end

########################################################
# D4M: Dynamic Distributed Dimensional Data Model
# Architect: Dr. Jeremy Kepner (kepner@ll.mit.edu)
# Software Engineer: Alexander Chen (alexc89@mit.edu)
########################################################


@symbol_func function q_95(cur_reactor::AbstractReactor)
  cur_q = 5.0

  cur_q *= cur_reactor.B_0

  cur_q *= a(cur_reactor) ^ 2

  cur_q *= _f_q(cur_reactor)

  cur_q /= cur_reactor.R_0

  cur_q /= cur_reactor.I_P

  cur_q
end


# SPDX-License-Identifier: MIT

"""
    struct CompiledFSM{T<:Semifield}
        π        # vector of initial probabilities
        ω        # vector of final probabilities
        A        # matrix of transition probabilities
        Aᵀ       # transpose of `A`
        pdfmap   # mapping state -> pdfindex
    end

Compiled FSM: matrix/vector format of an FSM used by inference
algorithms. Note that in this form, every state is an emitting state.
"""
struct CompiledFSM{T<:Semifield}
    π::AbstractVector{T}
    ω::AbstractVector{T}
    A::AbstractMatrix{T}
    Aᵀ::AbstractMatrix{T}
    pdfmap::AbstractVector{PdfIndex}
end

fsmtype(::CompiledFSM{T}) where T<:Semifield = T

"""
    compile(fsm; allocator = spzeros)

Compile `fsm` into a inference-friendly format [`CompiledFSM`](@ref).
`allocator` is a function analogous to `zeros` which creates an
N-dimensional array and fills it with zero elements.
"""
function compile(fsm::FSM{T}; allocator = spzeros) where T
    allstates = collect(states(fsm))
    S = length(allstates)

    # Initial states' probabilities.
    π = allocator(T, S)
    for s in filter(isinit, allstates) π[s.id] = s.initweight end

    # Final states' probabilities.
    ω = allocator(T, S)
    for s in filter(isfinal, allstates) ω[s.id] = s.finalweight end

    # Transition matrix.
    A = allocator(T, S, S)
    Aᵀ = allocator(T, S, S)
    for src in allstates
        for link in links(fsm, src)
            A[src.id, link.dest.id] = link.weight
            Aᵀ[link.dest.id, src.id] = link.weight
        end
    end

    pdfmap = [s.pdfindex for s in sort(allstates, by = p -> p.id)]

    CompiledFSM{T}(π, ω, A, Aᵀ, pdfmap)
end

"""
    gpu(cfsm)

Move the compiled fsm `cfsm` to GPU.
"""
function gpu(cfsm::CompiledFSM{T}) where T
    if ! issparse(cfsm.π)
        return CompiledFSM{T}(
            CuArray(cfsm.π),
            CuArray(cfsm.ω),
            CuArray(cfsm.A),
            CuArray(cfsm.Aᵀ),
            CuArray(cfsm.pdfmap)
        )
    end

    A = CuSparseMatrixCSC(cfsm.A)
    Aᵀ = CuSparseMatrixCSC(cfsm.Aᵀ)
    return CompiledFSM{T}(
        CuSparseVector(cfsm.π),
        CuSparseVector(cfsm.ω),
        CuSparseMatrixCSR(Aᵀ.colPtr, Aᵀ.rowVal, Aᵀ.nzVal, A.dims),
        CuSparseMatrixCSR(A.colPtr, A.rowVal, A.nzVal, A.dims),
        CuArray(cfsm.pdfmap)
    )
end

function Base.convert(::Type{CompiledFSM{NT}}, cfsm) where NT <: Semifield
    CompiledFSM(
        copyto!(similar(cfsm.π, NT), cfsm.π),
        copyto!(similar(cfsm.ω, NT), cfsm.ω),
        copyto!(similar(cfsm.A, NT), cfsm.A),
        copyto!(similar(cfsm.Aᵀ, NT), cfsm.Aᵀ),
        copyto!(similar(cfsm.pdfmap, Int), cfsm.pdfmap)
    )
end
Base.convert(::Type{T}, cfsm::T) where T <: CompiledFSM = cfsm


include("Graph.jl")
include("Centrality.jl")

fileName = string("Dataset/", ARGS[1], ".txt")
G0 = readGraph(fileName, "unweighted")
G = getLLC(G0)
lg = open("log.txt", "a")
println(lg, ARGS[1], " ", G0.n, " ", G0.m, " ", G.n, " ", G.m)

println(lg, "DC : ", getPc(Degree(G)))
println(lg, "BC : ", getPc(BetweennessCentrality(G)))
println(lg, "CC : ", getPc(ClosenessCentrality(G)))
println(lg, "EC : ", getPc(EigenvectorCentrality(G)))
println(lg, "PR : ", getPc(PageRank(G)))
println(lg, "IC : ", getPc(InformationCentrality(G)))
println(lg, "FDC : ", getPc(ForestDistanceClosenessCentrality(G)))
println(lg)


# Linear algebra functions

const σ0 = TensorMap(Matrix{ComplexF64}([1 0; 0 1]), ℂ^2, ℂ^2) # Identity matrix
const σ1 = TensorMap(Matrix{ComplexF64}([0 1; 1 0]), ℂ^2, ℂ^2) # Pauli X matrix
const σ2 = TensorMap(Matrix{ComplexF64}([0 -im; im 0]), ℂ^2, ℂ^2) # Pauli Y matrix
const σ3 = TensorMap(Matrix{ComplexF64}([1 0; 0 -1]), ℂ^2, ℂ^2) # Pauli Z matrix

const σp = 1 / 2 * (σ1 + im * σ2) # Pauli+ matrix
const σm = 1 / 2 * (σ1 - im * σ2) # Pauli- matrix

function LinearAlgebra.tr(t::PEPOTensor{S}) where {S<:IndexSpace}
    t = t.tensor
    @tensor tTr[1 2; 3 4] := t[1 2 x; 3 4 x]
end

function LinearAlgebra.ishermitian(p::PEPOTensor; phys_only=false)
    t = p.tensor
    phys_only || return LinearAlgebra.ishermitian(t)
    tT = permute(t, (1, 2, 6), (4, 5, 3))
    for (key, block) in blocks(t)
        block ≈ conj(blocks(tT)[key]) || return false
    end
    return true
end

function LinearAlgebra.vec(p::PEPOTensor)
    t = p.tensor
    return VectorPEPOTensor(_vec(t)...)
end

function _vec(t::TensorMap{S,3,3}) where {S<:EuclideanSpace}
    dom = TensorKit.domain(t)
    in = dom[3] # Physical index corresponding to the input
    cod = TensorKit.codomain(t)
    out = cod[3] # Physical index corresponding to the output
    tt = permute(t, (1, 2, 3, 6), (4, 5))
    u = unitary(fuse(out ⊗ in), out ⊗ flip(in))
    @tensor ttt[1 2 ph; 3 4] := u[ph; x1 x2] * tt[1 2 x1 x2; 3 4]
    return ttt, u
end

function unvec(p::VectorPEPOTensor)
    t = p.tensor
    u = p.unitary
    iu = inv(u)
    @tensor tt[1 2 x1 x2; 3 4] := iu[x1 x2; ph] * t[1 2 ph; 3 4]
    ttt = permute(tt, (1, 2, 3), (5, 6, 4))
    return PEPOTensor(ttt)
end

function conjugate(t::AbstractTensorMap)
    return TensorMap(conj(t.data), codomain(t), domain(t))
end

const cj = conjugate

@doc raw"""
    trace(t::PEPOTensor{S}, bond::S) where {S<:IndexSpace}

"""
function trace(t::PEPOTensor{S}, bond::S) where {S<:IndexSpace}
    @tensor pt[w s; n e] := t[w s x; n e x]
    guess = InfiniteMPS([codomain(pt)[2]], [bond])
    ts, _, _ = leading_boundary(guess, InfiniteMPO(pt), Vumps(; maxiter=400, verbose=true))
    return expectation_value(ts, pt, 1), ts
end

trace(t::PEPOTensor{ComplexSpace}, χ::Integer) = trace(t, ℂ^χ)


import ThermofluidQuantities: FrictionFactor

@nondimvar RoughnessCoefficient

function FrictionFactor(Re::ReynoldsNumber,eD::ET) where ET <: DimensionlessPhysicalQuantity
    xr = (5,Inf)
    xroot = find_zero(x -> -2.0*log10(eD/3.7 + 2.51/Re*x) - x,xr,order=16)
    FrictionFactor(1/xroot^2)
end


function shock_loading(m::RealBondMkup, TTT_jump::Matrix{Float64})
    exo = m.exogenous_shocks
    endo = m.endogenous_states

    _RRR = zeros(n_backward_looking_states(m), n_shocks_exogenous(m))
    RRR  = zeros(n_model_states(m), n_shocks_exogenous(m))

    # Technology shock loading on states
    _RRR[endo[:z′_t], exo[:z_sh]] .= 1.

    # Monetary policy shock loading on states
    _RRR[endo[:mon′_t], exo[:mon_sh]] .= 1.

    # Markup shock loading on states
    _RRR[endo[:mkp′_t], exo[:mkp_sh]] .= 1.

    # Loading on states and jumps
    RRR[1:n_backward_looking_states(m), :] = _RRR
    RRR[n_backward_looking_states(m)+1:end, :] = TTT_jump*_RRR

    return RRR
end


export getAcousticPointSource,loc2cs;

function getElasticRhs(Minv,TYPE)
## Generating the right hand side
if Minv.dim==3
	qx = zeros(TYPE,tuple(Minv.n + [1; 0; 0]...));
	qy = zeros(TYPE,tuple(Minv.n + [0; 1; 0]...));
	qz = zeros(TYPE,tuple(Minv.n + [0; 0; 1]...));
	src = [div(Minv.n[1],2),div(Minv.n[2],2),1];
	qz[src[1],src[2],1] = -2.0./(Minv.h[1]);
	qx[src[1],src[2],1] = -1.0./(Minv.h[3]);
	qx[src[1]+1,src[2],1] = +1.0./(Minv.h[3]);
	qy[src[1],src[2],1] = -1.0./(Minv.h[3]);
	qy[src[1],src[2]+1,1] = +1.0./(Minv.h[3]);
	q = vec([qx[:]; qy[:]; qz[:];]);
else
	qx = zeros(TYPE,tuple(Minv.n + [1; 0]...))
	qy = zeros(TYPE,tuple(Minv.n + [0; 1]...))
	qy[div(Minv.n[1],2),1] = -2.0./(Minv.h[1]);
	qx[div(Minv.n[1],2),1] = -1.0./(Minv.h[2]);
	qx[div(Minv.n[1],2)+1,1] = +1.0./(Minv.h[2]);
	# qy[div(Minv.n[1],2),div(Minv.n[2],2)] = -2./(Minv.h[1]);
	# qx[div(Minv.n[1],2),div(Minv.n[2],2)] = -1./(Minv.h[2]);
	# qx[div(Minv.n[1],2)+1,div(Minv.n[2],2)] = +1./(Minv.h[2]);
	q = vec([qx[:]; qy[:]])
end
return q;
end

function getTopPointSrc(Minv)
if Minv.dim==3
	src = [div(Minv.n[1]+1,2);div(Minv.n[2]+1,2);1];
else
    src = [div(Minv.n[1]+1,2);1]
end
return src;
end

function loc2cs(loc1::Int64,loc2::Int64,n::Array{Int64,1})
@inbounds cs = loc1 + (loc2-1)*n[1];
return cs;
end

function loc2cs3D(loc1::Int64,loc2::Int64,loc3::Int64,n::Array{Int64,1})
@inbounds cs = loc1 + (loc2-1)*n[1] + (loc3-1)*n[1]*n[2];
return cs;
end

function loc2cs(n::Array{Int64},sub::Array{Int64})
if length(sub)==2
	return loc2cs(sub[1],sub[2],n);
else
	return loc2cs3D(sub[1],sub[2],sub[3],n);
end
end

function getAcousticPointSource(Minv,TYPE,src = getTopPointSrc(Minv))
## Generating the right hand side
n_nodes = Minv.n .+ 1;
q = zeros(TYPE,tuple(n_nodes...));
q[loc2cs(n_nodes,src)] = 1.0./(norm(Minv.h)^2);
return q,src;
end


module LogDensities

using SparseQuadratureGrids, StaticArrays, ConstrainedParameters, DiffBase, ForwardDiff, Optim

import  Base: show, getindex, setindex!, size, IndexStyle, convert, length
import  ConstrainedParameters: type_length, param_type_length, construct, log_jacobian, update!
import  Optim:  Options,
                initial_state,
                update_state!,
                NewtonTrustRegionState,
                solve_tr_subproblem!,
                assess_convergence,
                value,
                gradient
#import  DiffBase: MutableDiffResult

export  Data,
        parameters,
        CovarianceMatrix,
        PositiveVector,
        ProbabilityVector,
        RealVector,
        Model,
        ModelParam,
        construct,
        log_density,
        log_jacobian!,
        quad_form,
        inv_det,
        inv_root_det,
        root_det,
        log_root_det,
        trace_inverse,
        lpdf_InverseWishart,
        lpdf_normal,
        logit,
        logistic,
        update!,
        type_length,
        ModelRank,
        StaticRank,
        DynamicRank,
        Full,
        FixedRank,
        LDR,
        Dynamic,
        MarginalBuffer,
        optimize!,
        ModelDiff

include("model_utilities.jl")
include("parameter_methods.jl")
include("model_derivatives.jl")
include("model.jl")

end # module


#=
dtft.jl
discrete-time Fourier transform
The main purpose of this code is for testing the NUFFT code.
2019-06-05, Jeff Fessler, University of Michigan
Note: the multi-dimensional version uses collect(Tuple(idx))
that might be inefficient, but for efficiency one should just use NUFFT.

see also MIRT/time/dtft.jl
=#

export dtft_init, dtft, dtft_adj

#using LinearAlgebra: norm
using Distributed: @sync, @distributed, pmap
using SharedArrays: SharedArray, sdata
using LinearMapsAA: LinearMapAA, LinearMapAO

"""
    X = dtft(w, x ; n_shift=?)
1D DTFT

``X[m] = \\sum_{n=0}^{N-1} x[n] \\exp(-i w[m] (n - n_{shift})), m=1,…,M``

in
- `w::AbstractVector{<:Real}`	`[M]` frequency locations ("units" radians/sample)
- `x::AbstractVector{<:Number}`	`[N]` 1D signal

option
- `n_shift::Real` often is N/2; default 0

out
- `X::AbstractVector{ComplexF64}` `[M]` DTFT
"""
function dtft(w::AbstractVector{<:Real}, x::AbstractVector{<:Number}
		; n_shift::Real = 0)
	return dtft_loop_n(w, x ; n_shift=n_shift)
end

"""
    X = dtft(w, x ; n_shift=?)
multi-dimensional DTFT (DSFT)

``X[m] = \\sum_{n=0}^{N-1} x[n] \\exp(-i w[m,:] (n - n_{shift})), m=1,…,M``
where here `n` is a `CartesianIndex`

in
- `w::AbstractMatrix{<:Real}`	`[M,D]` frequency locations ("units" radians/sample)
- `x::AbstractArray{<:Number}`	`[(Nd)]` D-dimensional signal

option
- `n_shift::AbstractVector{<:Real}` often is N/2; default zeros(D)

out
- `X::AbstractVector{ComplexF64}` `[M]` DTFT
"""
function dtft(w::AbstractMatrix{<:Real}, x::AbstractMatrix{<:Number}
	; n_shift::AbstractVector{<:Real} = zeros(Int, ndims(x)),
)
	return dtft_loop_n(w, x ; n_shift=n_shift)
end

"""
    d = dtft_init(w, N ; n_shift=?)
for 1D DTFT

in
- `w::AbstractVector{<:Real}`	`[M]` frequency locations ("units" radians/sample)
- `N::Int` size of signal `x`

option
- `n_shift::Real` often is N/2; default 0
- `T::DataType` default `ComplexF64` for testing NUFFT accuracy

out
- `d::NamedTuple` with fields
	`dtft = x -> dtft(x), adjoint = y -> dtft_adj(y), A=LinearMapAO`
"""
function dtft_init(
	w::AbstractVector{<:Real},
	N::Int ;
	n_shift::Real = 0,
	T::DataType = ComplexF64,
)

	M = length(w)
	forw = x -> dtft(w, x ; n_shift=n_shift)
	back = y -> dtft_adj(w, y, N ; n_shift=n_shift)
	A = LinearMapAA(forw, back, (M, N),
		(name="dtft1", N=(N,)) ; T = ComplexF64,
		operator = true,
	)
	return (dtft=forw, adjoint=back, A=A)
end

"""
    d = dtft_init(w, N ; n_shift=?)
for multi-dimensional DTFT (DSFT)

in
- `w::AbstractMatrix{<:Real}`	`[M,D]` frequency locations ("units" radians/sample)
- `N::Dims` `[D]` dimensions of signal `x`

option
- `n_shift::AbstractVector{<:Real}` often is N/2; default zeros(D)
- `T::DataType` default `ComplexF64` for testing NUFFT accuracy

out
- `d::NamedTuple` with fields
	`dtft = x -> dtft(x), adjoint = y -> dtft_adj(y), A=LinearMapAO`
"""
function dtft_init(
	w::AbstractMatrix{<:Real},
	N::Dims ;
	n_shift::AbstractVector{<:Real} = zeros(Int, length(N)),
)

	D = size(w,2)
	length(N) != D && throw(DimensionMismatch("length(N) vs D=$D"))
	length(n_shift) != D && throw(DimensionMismatch("length(n_shift) vs D=$D"))
	M = size(w,1)
	forw = x -> dtft(w, x ; n_shift=n_shift)
	back = y -> dtft_adj(w, y, N ; n_shift=n_shift)
	A = LinearMapAA(forw, back, (M, prod(N)),
		(name="dtft$(length(N))", N=N) ; T = ComplexF64,
		operator = true, idim = N,
	)
	return (dtft=forw, adjoint=back, A=A)
end

"""
    x = dtft_adj(w, X, N ; n_shift=?)
adjoint for 1D DTFT

``x[n] = \\sum_{m=1}^M X[m] \\exp(i w[m] (n - n_{shift})), n=0,…,N-1``

This is the *adjoint* (transpose) of `dtft`, not an *inverse* DTFT.

in
- `w::AbstractVector{<:Real}`	`[M]` frequency locations ("units" radians/sample)
- `X::AbstractVector{ComplexF64}`	`[M]` spectrum values
- `N::Int` size of signal `x`

option
- `n_shift::Real` often is N/2; default 0

out
- `x::AbstractVector{<:Number}`	signal [N]
"""
function dtft_adj(w::AbstractVector{<:Real}, X::AbstractVector{<:Number},
	N::Int ; n_shift::Real = 0,
)
	M = length(w)
	out = Array{ComplexF64}(undef, N)
	nshift1 = n_shift + 1
	X = ComplexF64.(X)
	for n=1:N # loop over output signal samples
		out[n] = sum(cis.((n - nshift1) * w) .* X)
	end
	return out
end

"""
    x = dtft_adj(w, X, N ; n_shift=?)
adjoint for multi-dimensional DTFT (DSFT)

``x[n] = \\sum_{m=1}^M X[m] \\exp(i w[m,:] (n - n_{shift})), n=0,…,N-1``
where here `n` is a `CartesianIndex`

in
- `X::AbstractVector{ComplexF64}` `[M]` DTFT
- `w::AbstractMatrix{<:Real}` `[M,D]` frequency locations ("units" radians/sample)
- `N::Dims`	`[D]` dimensions of signal `x`

option
- `n_shift::AbstractVector{<:Real}` often is `N/2`; default `zeros(D)`

out
- `x::AbstractArray{<:Number}` `[(N)]` `D`-dimensional signal
"""
function dtft_adj(
	w::AbstractMatrix{<:Real},
	X::AbstractVector{<:Number},
	N::Dims ;
	n_shift::AbstractVector{<:Real} = zeros(Int, ndims(x)),
)

	M,D = size(w)
	nshift1 = n_shift .+ 1

	T = (eltype(w) == Float64) ? ComplexF64 : ComplexF32
	out = Array{T}(undef, N)

	for n=1:prod(N)
		idx = CartesianIndices(N)[n]
		tmp = cis.(w * (collect(Tuple(idx)) - nshift1))
		tmp = transpose(cis.(w * (collect(Tuple(idx)) - nshift1))) * X
		out[idx] = cis.(-w * (collect(Tuple(idx)) - nshift1))' * X
	end
	return out
end

# 1D loop over signal values
function dtft_loop_n(w::AbstractVector{<:Real}, x::AbstractVector{<:Number}
	; n_shift::Real = 0,
)
	N = length(x)
	out = fill(ComplexF64(0), size(w))
	nshift1 = n_shift + 1
	for n=1:N # tried @simd but not faster
		out .+= x[n] * cis.(-(n-nshift1) * w)
	end
	return out
end

# dD loop over signal values
# * `w [M,D]`
# * `x [(Nd)]`
function dtft_loop_n(
	w::AbstractMatrix{<:Real},
	x::AbstractMatrix{<:Number} ;
	n_shift::AbstractVector{<:Real} = zeros(Int, size(w,2)),
)
	M,D = size(w)
	N = size(x)
	length(N) != D && throw(DimensionMismatch("length(N) vs D=$D"))
	length(n_shift) != D && throw(DimensionMismatch("length(n_shift) vs D=$D"))
	ndims(x) != D && throw(DimensionMismatch("ndims(x) vs D=$D"))
	nshift1 = n_shift .+ 1
	T = (eltype(w) == Float64) ? ComplexF64 : ComplexF32
	out = fill(T(0), M)
	for n=1:prod(N)
		idx = CartesianIndices(N)[n]
		out .+= x[idx] * cis.(-w * (collect(Tuple(idx)) - nshift1))
	end
	return out
end

# 1D loop over frequencies
function dtft_loop_m(
	w::AbstractVector{<:Real},
	x::AbstractVector{<:Number} ;
	n_shift::Real = 0,
)
	M = length(w)
	N = length(x)
	out = Array{ComplexF64}(undef, M)
	nn = -((0:(N-1)) .- n_shift)
	x = ComplexF64.(x)
	for m=1:M
		out[m] = sum(cis.(nn * w[m]) .* x)
	end
	return out
end

# 1D matrix * vector
function dtft_matvec(
	w::AbstractVector{<:Real},
	x::AbstractVector{<:Number} ;
	n_shift::Real = 0,
)
	N = length(x)
	return cis.(w * (-((0:(N-1)) .- n_shift))') * x
end

# 1D single frequency
function dtft_one_w(
	w::Real,
	x::AbstractVector{<:Number} ;
	n_shift::Real = 0,
)
	N = length(x)
	return sum(cis.(-((0:(N-1)) .- n_shift) * w) .* x)
end

# 1D pmap
function dtft_pmap_m(
	w::AbstractVector{<:Real},
	x::AbstractVector{<:Number}
	; n_shift::Real = 0,
)
	tmp = w -> dtft_one_w(w, x ; n_shift=n_shift)
	return pmap(tmp, w)
end

# 1D distributed loop over frequencies
function dtft_dist_m(
	w::AbstractVector{<:Real},
	x::AbstractVector{<:Number} ;
	n_shift::Real = 0,
)
	M = length(w)
	N = length(x)
	out = SharedArray{ComplexF64}(M)
	nn = -((0:(N-1)) .- n_shift)
	x = ComplexF64.(x)
	@sync @distributed for m=1:M
		out[m] = sum(cis.(nn * w[m]) .* x)
	end
	return sdata(out)
end


function generateU(ul::Vector{Complex128},xn::Vector{Complex128})
	L = length(ul) #Num of basis functions
	N = length(xn) #signal length
	p = 1 #calculate U matrices up to this num (1) for U0,U1
	K = ifloor((N-1-p)/2)
    NPOW = 3

	ul_inv = [1/ul[i] for i=1:L]
	ul_invK = [Base.power_by_squaring(ul_inv[i],K) for i=1:L]

	#initialize memory
	g0 = zeros(Complex128,L)
	g0_K = zeros(Complex128,L)
    g0_NK = zeros(Complex128,L)
	D0 = zeros(Complex128,L)
	U0 = zeros(Complex128,(L,L))
	U1 = zeros(Complex128,(L,L))
    ul_invk = ones(Complex128,L)
	for i=0:(K)
		for j=0:(L-1)
			g0[j+1] += ul_invk[j+1]*xn[i+1]
			g0_K[j+1] += ul_invk[j+1]*xn[K+1+i+1]
            g0_NK[j+1] += ul_invk[j+1]*xn[N-K+i]
            D0[j+1] += (i + 1)*xn[i+1]*ul_invk[j+1] + (K - i)*xn[i + K + 1+1]*ul_invk[j+1]*ul_inv[j+1]*ul_invK[j+1]
            ul_invk[j+1] = ((i%NPOW)==(NPOW-1)) ? Base.power_by_squaring(ul_inv[j+1],i+1) : ul_invk[j+1]*ul_inv[j+1]
		end
	end

	for i=1:L
		for j=1:i
			U0[i,j] = (1/(ul[i]-ul[j]))*(ul[i]*g0[j] - ul[j]*g0[i]
					+ ul_invK[j]*g0_K[i] - ul_invK[i]*g0_K[j])
			U0[j,i] = U0[i,j]
			U1[i,j] = (1/(ul[i]-ul[j]))*(ul[i]*ul[j]*g0[j] - ul[j]*ul[i]*g0[i]
					+ ul_invK[j]*ul[i]*g0_K[i] - ul_invK[i]*ul[j]*g0_K[j])
			U1[j,i] = U1[i,j]
		end
		U0[i,i] = D0[i]
		U1[i,i] = D0[i]*ul[i]-ul[i]*g0[i]+ul_invK[i]*g0_K[i]
	end
	return (U0,U1,g0,g0_K, g0_NK)
end

function xft(xn::Vector{Complex128},ul::Vector{Complex128},Z::Vector{Complex128},q::Float64)

    N = length(xn)
    if length(Z)<N
        error("Length of Z must be at least the length of N")
    end
    M = ifloor((N-2)/2)+1
    (U0,U1,g0,g0_K,g0_NK) = generateU(ul,xn)
    genR(z,U0,U1) = U0-U1/z
    U0dagU0 = ctranspose(U0)*U0
    U1dagU1 = ctranspose(U1)*U1
    U1dagU0 = ctranspose(U1)*U0
    U0dagU1 = ctranspose(U0)*U1
    genRdagG0(R,g0) = ctranspose(R)*g0
    genRdagR(z,U0dagU0,U1dagU1,U1dagU0,U0dagU1) = U0dagU0 + U1dagU1 -z*U1dagU0 - U0dagU1/z
    Spec = zeros(Complex128,length(Z))
    for k=1:length(Z)
        R  = genR(Z[k],U0,U1)
        RC2 = genRdagG0(R,g0_NK*Z[k]^(-N+M))
        RR = genRdagR(Z[k],U0dagU0,U1dagU1,U1dagU0,U0dagU1)
        Spec[k] = (transpose(g0_K*Z[k]^(-M))*(\((RR+q^2),RC2)))[1]
    end
    return Spec+fft(vcat(xn,zeros(length(Z)-N)))
end

function rrt(xn::Vector{Complex128},ul::Vector{Complex128},Z::Vector{Complex128},q::Float64)
    N = length(xn)
    M = ifloor((N-2)/2)+1
    (U0,U1,g0,g0_K,g0_NK) = generateU(ul,xn)
    genR(z,U0,U1) = U0-U1/z
    U0dagU0 = ctranspose(U0)*U0
    U1dagU1 = ctranspose(U1)*U1
    U1dagU0 = ctranspose(U1)*U0
    U0dagU1 = ctranspose(U0)*U1
    genRdagG0(R,g0) = ctranspose(R)*g0
    genRdagR(z,U0dagU0,U1dagU1,U1dagU0,U0dagU1) = U0dagU0 + U1dagU1 -z*U1dagU0 - U0dagU1/z
    Spec = zeros(Complex128,length(Z))
    for k=1:length(Z)
        R  = genR(Z[k],U0,U1)
        RC = genRdagG0(R,g0)
        RR = genRdagR(Z[k],U0dagU0,U1dagU1,U1dagU0,U0dagU1)
        Spec[k] = (transpose(g0)*(\((RR+q^2),RC)))[1]
    end
    return Spec
end

function xczt(xn::Vector{Complex128},A::Complex128,W::Complex128,K::Int,q::Float64)
    N = length(xn)
    Q = N*sum(abs(xn))
    q = Q*q
    Nul = minimum((N,100))
    ul = [exp(im*2*pi/Nul)^(i-1) for i=1:Nul] #fourier basis for expanding Krylov operators U0,U1
    Z = [A*W^(-(i-1)) for i=1:K] #this is the spiral the CZT is to be evaluated on
    M = ifloor((N-2)/2)+1
    (U0,U1,g0,g0_K,g0_NK) = generateU(ul,xn)
    genR(z,U0,U1) = U0-U1/z
    U0dagU0 = ctranspose(U0)*U0
    U1dagU1 = ctranspose(U1)*U1
    U1dagU0 = ctranspose(U1)*U0
    U0dagU1 = ctranspose(U0)*U1
    genRdagG0(R,g0) = ctranspose(R)*g0
    genRdagR(z,U0dagU0,U1dagU1,U1dagU0,U0dagU1) = U0dagU0 + U1dagU1 -z*U1dagU0 - U0dagU1/z
    Spec = zeros(Complex128,length(Z))
    for k=1:length(Z)
        R  = genR(Z[k],U0,U1)
        RC2 = genRdagG0(R,g0_NK*Z[k]^(-N+M))
        RR = genRdagR(Z[k],U0dagU0,U1dagU1,U1dagU0,U0dagU1)
        Spec[k] = (transpose(g0_K*Z[k]^(-M))*(\((RR+q^2),RC2)))[1]
    end
    return Spec+czt(A,W,K,xn)
end

function rczt(xn::Vector{Complex128},A::Complex128,W::Complex128,K::Int,q::Float64)
    N = length(xn)
    Q = N*sum(abs(xn))
    q = Q*q
    Nul = minimum((N,100))
    ul = [exp(im*2*pi/Nul)^(i-1) for i=1:Nul] #fourier basis for expanding Krylov operators U0,U1
    Z = [A*W^(-(i-1)) for i=1:K] #this is the spiral the CZT is to be evaluated on
    M = ifloor((N-2)/2)+1
    (U0,U1,g0,g0_K,g0_NK) = generateU(ul,xn)
    genR(z,U0,U1) = U0-U1/z
    U0dagU0 = ctranspose(U0)*U0
    U1dagU1 = ctranspose(U1)*U1
    U1dagU0 = ctranspose(U1)*U0
    U0dagU1 = ctranspose(U0)*U1
    genRdagG0(R,g0) = ctranspose(R)*g0
    genRdagR(z,U0dagU0,U1dagU1,U1dagU0,U0dagU1) = U0dagU0 + U1dagU1 -z*U1dagU0 - U0dagU1/z
    Spec = zeros(Complex128,length(Z))
    for k=1:length(Z)
        R  = genR(Z[k],U0,U1)
        RC = genRdagG0(R,g0)
        RR = genRdagR(Z[k],U0dagU0,U1dagU1,U1dagU0,U0dagU1)
        Spec[k] = (transpose(g0)*(\((RR+q^2),RC)))[1]
    end
    return Spec
end


using DRIPs
using Test

@testset "Sims2010.jl" begin
    # Write your own tests here.
    println("Running tests ...")
	    β = 0.9;
		ω = 1.0;
		A = [0.95 0.0; 0.0 0.4];
		Q = [√0.0975 0.0; 0.0 √0.86];
		H = [1.0; 1.0];
	println("Solving Sims (2011) ...")
	    p   = Drip(ω,β,A,Q,H)
	    @test p.ss.err < 1e-4
    println("Solving Sims (2011) with fixed capacity ...")
	    cap = DRIPs.capacity(p)
	    capn= DRIPs.capacity(p; unit = "nat")
	    pf  = Drip(cap,β,A,Q,H;fcap=true)
	    @test pf.ss.err < 1e-4
	    @test pf.ss.Σ_p ≈ p.ss.Σ_p atol = 1e-3
    println("Solving transition dynamics for Sims (2011) ...")
	    S   = DRIPs.Signal([1 0; 0 1],[0 0; 0 0])
	    pt  = Trip(p,S;T=4000)
	    @test pt.err < 1e-4;
		@test pt.con_err < 1e-3;
    println("Checking IRFs ...")
		pirfs     = irfs(p,T = 15)
		psims1    = simulate(p,T = 500, burn = 100, seed = 0)
		psims2    = simulate(p,T = 500, burn = 100, N=100, seed = 0)
		@test psims1.x == psims2.x
		psims     = simulate(p,T = 500, burn = 100, N=1, seed = 0)
		ptss      = Trip(p,p.ss.Σ_1;T = 15)
		ptssirfs  = irfs(ptss,T = 15)
		Sp        = DRIPs.Signal([0 0; 0 0], [1 0; 0 1]);
		ptssirfsp = irfs(p,Sp, T = 30)
		ptssirfsp = irfs(p,Sp, T = 10)
		ptssirfsp = irfs(p,Sp, T = 15)
		ptssirfspn= irfs(p,Sp; reoptimize = false, T = 15)
		@test pirfs.x_hat ≈ ptssirfs.x_hat atol = 1e-3
		@test pirfs.x_hat ≈ ptssirfsp.x_hat atol = 1e-3
		@test pirfs.x_hat ≈ ptssirfspn.x_hat atol = 1e-3
	println("Final tests on internal functions ...")
		x(j) = 0.5^j
		sum  = DRIPs.infinitesum(x)
		@test sum ≈ 2 atol = 1e-4
end


function subgradient!(dd_model::DDModel)
    #get an upper bound
    alpha = dd_model.solver_data.alpha
    lower_bound = dd_model.current_lower_bound

    if dd_model.upper_bound_function != nothing
        dd_model.upper_bound = dd_model.upper_bound_function(dd_model)
        step = alpha*abs(lower_bound - dd_model.upper_bound)/(norm([dd_model.residual_equality;dd_model.residual_inequality])^2)
    else
        upper_bound = nothing
        step = alpha
    end

    lambda_equality_delta = step*dd_model.residuals_equality  #update multipliers
    lambda_inequality_delta = step*dd_model.residuals_inequality
    lambda_linkvar_delta = step*dd_model.residuals_linkvars

    return lambda_equality_delta,lambda_inequality_delta,lambda_linkvar_delta
end

#TODO
function cutting_plane!(dd_model::DDModel)
end


@testset "Systems" begin

    @testset "SPSystem" begin
        fs = equations(katsura(5))

        F = SPSystem(fs)

        @test size(F) == (6, 6)
        @test length(F) == 6
        x = rand(Complex{Float64}, 6)
        system_cache = cache(F, x)
        @test system_cache isa HomotopyContinuation.SystemNullCache
        u = zeros(Complex{Float64}, 6)
        evaluate!(u, F, x, system_cache)
        @test evaluate(F, x, system_cache) ≈ u
        @test evaluate(F, x) ≈ u
        @test F(x) ≈ u

        U = zeros(Complex{Float64}, 6, 6)
        jacobian!(U, F, x, system_cache)
        @test jacobian(F, x, system_cache) ≈ U
        @test jacobian(F, x) ≈ U

        evaluate_and_jacobian!(u, U, F, x, system_cache)
        @test all(evaluate_and_jacobian(F, x, system_cache) .≈ (u, U))

        @polyvar x y a b
        F = SPSystem(
            [x^2 + 2y * x + a * x + b + 3.0, b * x * y^3 + a + 2];
            parameters = [a, b],
        )
        @test size(F) == (2, 2)
        @test length(F) == 2
        x = rand(Complex{Float64}, 2)
        p = rand(Complex{Float64}, 2)
        system_cache = cache(F, x, p)
        @test system_cache isa HomotopyContinuation.SystemNullCache
        u = zeros(Complex{Float64}, 2)
        evaluate!(u, F, x, p, system_cache)
        @test evaluate(F, x, p, system_cache) ≈ u
        @test evaluate(F, x, p) ≈ u
        @test F(x, p) ≈ u

        U = zeros(Complex{Float64}, 2, 2)
        jacobian!(U, F, x, p, system_cache)
        @test jacobian(F, x, p, system_cache) ≈ U
        @test jacobian(F, x, p) ≈ U

        evaluate_and_jacobian!(u, U, F, x, p, system_cache)
        @test all(evaluate_and_jacobian(F, x, p, system_cache) .≈ (u, U))
        @test all(evaluate_and_jacobian(F, x, p) .≈ (u, U))
    end

    @testset "FixedHomotopy" begin
        f = SPSystem(equations(katsura(5)))
        g = SPSystem(equations(cyclic(6)))
        H = StraightLineHomotopy(f, g)

        F = FixedHomotopy(H, 0.23)
        x = rand(Complex{Float64}, 6)
        u = zeros(Complex{Float64}, 6)
        U = zeros(Complex{Float64}, 6, 6)

        system_cache = cache(F, x)
        @test system_cache isa HomotopyContinuation.FixedHomotopyCache
        u = zeros(Complex{Float64}, 6)
        evaluate!(u, F, x, system_cache)
        @test evaluate(F, x, system_cache) ≈ u

        U = zeros(Complex{Float64}, 6, 6)
        jacobian!(U, F, x, system_cache)
        @test jacobian(F, x, system_cache) ≈ U

        evaluate_and_jacobian!(u, U, F, x, system_cache)
        @test all(evaluate_and_jacobian(F, x, system_cache) .≈ (u, U))
    end

    @testset "FPSystem" begin
        F = FPSystem(equations(katsura(5)))
        x = rand(Complex{Float64}, 6)
        u = zeros(Complex{Float64}, 6)
        U = zeros(Complex{Float64}, 6, 6)

        system_cache = cache(F, x)
        @test system_cache isa HomotopyContinuation.FPSystemCache
        u = zeros(Complex{Float64}, 6)
        evaluate!(u, F, x, system_cache)
        @test evaluate(F, x, system_cache) ≈ u

        U = zeros(Complex{Float64}, 6, 6)
        jacobian!(U, F, x, system_cache)
        @test jacobian(F, x, system_cache) ≈ U

        evaluate_and_jacobian!(u, U, F, x, system_cache)
        @test all(evaluate_and_jacobian(F, x, system_cache) .≈ (u, U))
    end

    @testset "TotalDegreeSystem" begin
        @polyvar x y z

        F = SPSystem([x^4 - z^4, y^3 - z^3])
        w = ProjectiveVectors.PVector(rand(ComplexF64, 3))

        G = TotalDegreeSystem([x^4 - z^4, y^3 - z^3])
        @test size(G) == (2, 3)
        @test evaluate(F, w) ≈ evaluate(G, w)
        @test jacobian(F, w) ≈ jacobian(G, w)
        u, U = evaluate_and_jacobian(F, w)
        @test u ≈ evaluate(F, w)
        @test U ≈ jacobian(G, w)

        @polyvar x y z

        F_affine = SPSystem([x^4 - 1, y^3 - 1])
        w = rand(ComplexF64, 3)

        G_affine = TotalDegreeSystem([x^4 - 1, y^3 - 1]; affine = true)
        @test size(G_affine) == (2, 2)
        @test evaluate(F_affine, w) ≈ evaluate(G_affine, w)
        @test jacobian(F_affine, w) ≈ jacobian(G_affine, w)

        u, U = evaluate_and_jacobian(F_affine, w)
        @test u ≈ evaluate(F_affine, w)
        @test U ≈ jacobian(G_affine, w)
    end

    @testset "FixedParameterSystem" begin
        @polyvar x y a b
        f = SPSystem([x^2 + y - a, b * x + y * a^2]; parameters = [a, b])
        F = FixedParameterSystem(f, [2.0, 3.0])

        @test size(F) == (2, 2)
        @test length(F) == 2
        x = rand(Complex{Float64}, 2)
        system_cache = cache(F, x)
        @test system_cache isa HomotopyContinuation.FixedParameterSystemCache
        u = zeros(Complex{Float64}, 2)
        evaluate!(u, F, x, system_cache)
        @test evaluate(F, x, system_cache) ≈ u
        @test evaluate(F, x) ≈ u

        U = zeros(Complex{Float64}, 2, 2)
        jacobian!(U, F, x, system_cache)
        @test jacobian(F, x, system_cache) ≈ U
        @test jacobian(F, x) ≈ U

        evaluate_and_jacobian!(u, U, F, x, system_cache)
        @test all(evaluate_and_jacobian(F, x, system_cache) .≈ (u, U))

        set_parameters!(F, (4.0, 2.3))
        @test F.p == [4.0, 2.3]
    end

    @testset "SquaredUpSystem" begin
        # affine
        @polyvar x y z

        F = SPSystem([x^4 - 1, y^3 - 1, z^2 + x, x + y + z - 1, x^2 + z^2 - 3])
        w = rand(ComplexF64, 3)
        A = randn(ComplexF64, 3, 2)
        S = SquaredUpSystem(F, A)
        system_cache = cache(S, w)
        @test system_cache isa AbstractSystemCache

        @test size(S) == (3, 3)
        @test evaluate(S, w, system_cache) ≈ [LinearAlgebra.I A] * evaluate(F, w) atol = 1e-12
        @test jacobian(S, w, system_cache) ≈ [LinearAlgebra.I A] * jacobian(F, w) atol = 1e-12
        u, U = evaluate_and_jacobian(S, w, system_cache)
        @test u ≈ evaluate(S, w)
        @test U ≈ jacobian(S, w)

        # projective
        @polyvar x y z w

        A = randn(ComplexF64, 3, 2)
        F = [x^4 - 1, y^3 - 1, z^2 + x, x^2 + z^2 - 3, x + y + z - 1]
        G = SPSystem(homogenize([LinearAlgebra.I A] * F, w))
        F = SPSystem(homogenize(F, w))
        v = ProjectiveVectors.PVector(rand(ComplexF64, 4))

        S = SquaredUpSystem(F, A, [4, 3, 2, 2, 1])
        system_cache = cache(S, v)
        @test system_cache isa AbstractSystemCache

        @test size(S) == (3, 4)
        @test evaluate(S, v, system_cache) ≈ evaluate(G, v) atol = 1e-12
        @test jacobian(S, v, system_cache) ≈ jacobian(G, v) atol = 1e-12
        u, U = evaluate_and_jacobian(S, v, system_cache)
        @test u ≈ evaluate(S, v)
        @test U ≈ jacobian(S, v)

        @test HC.compute_numerically_degrees(S) == [4, 3, 2]
    end
end


module RecurrentTests

using Photon, Test

resetContext()
ctx = getContext()

gethistoric(steps=100, features=10) = KorA(randn(ctx.dtype,features,steps,4))

function basic_layer(T)

    d = gethistoric(100)

    l = T(20)
    pred = l(d)
    @test size(pred) == (20,4)

    l = T(20,2)
    pred = l(d)
    @test size(pred) == (20,4)

    l = T(20;dropout=0.5, bidirectional=true)
    pred = l(d)
    @test size(pred) == (20,4)

    l = T(20;dropout=0.5, last_only=false)
    pred = l(d)
    @test size(pred) == (20,100,4)

    l = T(20; initw=zeros, initb=zeros, initf=zeros)
    pred = l(d)
    @test size(pred) == (20,4)
    @test sum(pred) == 0.0

end

function lstm_model()

    model = Sequential(
        LSTM(20),
        Dense(100, relu),
        Dense(10)
    )

    pred = model(gethistoric(100))
    @test size(pred) == (10,4)

    pred = model(gethistoric(200))
    @test size(pred) == (10,4)

end

function lstm_model2()

    model = Sequential(
        LSTM(20, last_only=false),
        Dense(100, relu),
        Dense(10)
    )

    pred = model(gethistoric(100))
    @test size(pred) == (10,4)

    @test_throws DimensionMismatch model(gethistoric(200))
end

function gru_model()

    model = Sequential(
        GRU(20,2;dropout=0.5),
        Dense(100, relu),
        Dense(10)
    )

    pred = model(gethistoric(100))
    @test size(pred) == (10,4)
end

function rnn_model()

    model = Sequential(
        RNN(20,2;activation=:relu, dropout=0.5),
        Dense(100, relu),
        Dense(10)
    )

    pred = model(gethistoric(100))
    @test size(pred) == (10,4)
end

@testset "Recurrent" begin
    basic_layer(LSTM)
    basic_layer(GRU)
    basic_layer(RNN)
    lstm_model()
    lstm_model2()
    gru_model()
    rnn_model()
end

end


using Interpolation

using Test: @test, @testset

@testset "Interpolation Tests" begin
    @testset "Simple Analytic Examples" begin
        @test isapprox(polyint([1.0, 2.0], [1.0, 2.0])(3.0), 3.0)
        @test isapprox(polyint([1.0, 2.0, 3.0], [1.0, 4.0, 9.0])(4.0), 16.0)
        @test isapprox([polyint([1.0, 2.0, 3.0], [1.0, 4.0, 9.0], 1)(4.0)...], [10.0, 14.0])
    end

    @testset "Random Third Order Polynomials" begin
        for i in 1:10
            a,b,c,d = randn(4) # Roots

            xs = randn(4)

            ys = Float64[]
            for x in xs
                push!(ys, d*(x-a)*(x-b)*(x-c))
            end

            x = randn()
            poly = d*(x-a)*(x-b)*(x-c)

            @test isapprox(polyint(xs, ys)(x), poly)
        end
    end
end


cluster = [1 1 1 1; 2 2 3 3; 4 4 5 5]
X = Array(reshape(1:24, 2, 3, 4))

@testset "pool" begin
    @testset "GlobalPool" begin
        glb_cltr = [1 1 1 1; 1 1 1 1; 1 1 1 1]
        p = GlobalPool(:add, 3, 4)
        @test p(X) == sumpool(glb_cltr, X)
    end

    @testset "LocalPool" begin
        p = LocalPool(:add, cluster)
        @test p(X) == sumpool(cluster, X)
    end

    @testset "TopKPool" begin
        N = 10
        k, in_channel = 4, 7
        X = rand(in_channel, N)
        for T = [Bool, Float64]
            adj = rand(T, N, N)
            p = TopKPool(adj, k, in_channel)
            @test eltype(p.p) === Float32
            @test size(p.p) == (in_channel,)
            @test eltype(p.Ã) === T
            @test size(p.Ã) == (k, k)
            y = p(X)
            @test size(y) == (in_channel, k)
        end
    end
end


# Bernoulli Arm Model

type Bernoulli <: BanditArmBase
    armDist::Distributions.Bernoulli

    function Bernoulli( p::Real )
        new( Distributions.Bernoulli(p) )
    end
end

function pull!( arm::Bernoulli )
    return Distributions.rand(arm.armDist)
end

function tick!( arm::Bernoulli )
    # Do nothing
    nothing
end

function reset!( arm::Bernoulli )
    # Do nothing
    nothing
end


module TestGradientAlgebra
using Test
using LinearAlgebra
using WoodburyFactorizations
using BlockFactorizations
using CovarianceFunctions
using CovarianceFunctions: EQ, RQ, Dot, ExponentialDot, NN, GradientKernel,
        ValueGradientKernel, DerivativeKernel, ValueDerivativeKernel, input_trait,
        LazyMatrixSum, DotProductGradientKernelElement, SeparableProduct
using ForwardDiff
const AbstractMatOrFac = Union{AbstractMatrix, Factorization}

@testset "input trait property" begin
    # dot product
    k = CovarianceFunctions.Dot()^5
    h = CovarianceFunctions.Dot()^2 + 1 # quadratic
    @test input_trait(k+h) == DotProductInput()
    @test input_trait(k*h) == DotProductInput()

    # isotropic
    k = CovarianceFunctions.EQ()
    h = CovarianceFunctions.RQ(1.0) # quadratic
    @test input_trait(k+h) == IsotropicInput()
    @test input_trait(k*h) == IsotropicInput()

    # heterogeneous
    k = CovarianceFunctions.EQ()
    h = CovarianceFunctions.Dot()^2 # quadratic
    @test input_trait(k+h) == GenericInput()
    @test input_trait(k*h) == GenericInput()
end

@testset "gradient algebra" begin
    # test data
    d, n = 3, 7
    X = randn(d, n)

    # sum of homogeneous kernel types
    k = CovarianceFunctions.Dot()^5
    h = CovarianceFunctions.Dot()^2 + 1 # quadratic
    gk, gh = GradientKernel(k), GradientKernel(h)
    G = GradientKernel(k + h)
    K, Kk, Kh = gramian(G, X), gramian(gk, X), gramian(gh, X)
    MK, Mkk, Mkh = Matrix(K), Matrix(Kk), Matrix(Kh)
    @test K isa BlockFactorization
    @test K.A[1, 1] isa DotProductGradientKernelElement # this means it correctly consolidated the sum kernel into one dot product kernel
    @test MK ≈ Mkk + Mkh

    # sum of heterogeneous kernel types
    eq = CovarianceFunctions.EQ() # EQ
    line = CovarianceFunctions.Dot() # dot
    quad = line^2 + 1
    cosine = CovarianceFunctions.Cosine(randn(d))
    x, y = X[:, 1], X[:, 2]
    # TODO: Chained, VerticalScaling, Warped, SeparableSum
    heterogeneous_combinations = [eq + line, eq + quad, eq * cosine, eq * quad,
                                SeparableProduct(k, h, k), line * line]
    gradient_kernels = [GradientKernel, ValueGradientKernel]
    for gradient_kernel in gradient_kernels
        for kh in heterogeneous_combinations
            # testing product gradient kernel
            G = gradient_kernel(kh)
            kh_control = (x, y) -> (kh)(x, y)
            G_control = gradient_kernel(kh_control)
            Gxy = G(x, y)
            Gxy_control = G_control(x, y)
            @test typeof(Gxy) != typeof(Gxy_control) # means special structure was discovered
            @test Matrix(Gxy) ≈ Gxy_control
        end
    end
end

end


using Test
using CompScienceMeshes

for T in [Float32, Float64]
    mesh = meshrectangle(T(1.0), T(1.0), T(1.0))
    faces = skeleton(mesh, 2)
    local verts = vertices(mesh, faces.faces[1])
    p = simplex(verts)

    @test p.vertices == [
        point(T, 0.0, 0.0, 0.0),
        point(T, 0.0, 1.0, 0.0),
        point(T, 1.0, 0.0, 0.0)]
    @test p.tangents == [
        point(T, -1.0, 0.0, 0.0),
        point(T, -1.0, 1.0, 0.0)]
    @test p.normals == [
        point(T, 0.0, 0.0, -1.0)]
    @test p.volume == T(0.5)
end


# this file tests various aspects of the advection routines
using MagmaThermoKinematics
using ParallelStencil
using ParallelStencil.FiniteDifferences3D
using Plots
using LinearAlgebra
using SpecialFunctions
using Test

const CreatePlots = false      # easy way to deactivate plotting throughout

# Initialize for multiple threads (GPU is not tested here)
@init_parallel_stencil(Threads, Float64, 3);    # initialize parallel stencil in 3D

function test_Interpolation(Dimension="2D", InterpolationMethod="Linear")
  # test interpolation methods in 2D and 3D

  if Dimension=="2D"
    # Model parameters
    W,H                     =   1.,  1.;                                    # Width, Length, Height

    # Define coarse grid
    Nx, Nz                  =   33, 33;                                     # resolution of coarse grid
    dx,dz                   =   W/(Nx-1), H/(Nz-1);                         # grid size [m]
    x,z                     =   0:dx:((Nx-1)*dx), 0:dz:((Nz-1)*dz);         # 1D coordinate arrays
    coords                  =   collect(Iterators.product(x,z))             # generate coordinates from 1D coordinate vectors
    X,Z                     =   (x->x[1]).(coords), (x->x[2]).(coords);     # transfer coords to 3D arrays
    Grid, Spacing           =   (x,z), (dx,dz);

    # Define fine grid
    Nx_f, Nz_f              =   65, 65;                                     # resolution of fine grid
    dx_f,dz_f               =   W/(Nx_f-1), H/(Nz_f-1);                     # grid size [m]
    x_f,z_f                 =   0:dx_f:H, 0:dz_f:H;                         # 1D coordinate arrays
    coords_f                =   collect(Iterators.product(x_f,z_f))         # generate coordinates from 1D coordinate vectors
    X_f,Z_f                 =   (x->x[1]).(coords_f), (x->x[2]).(coords_f); # transfer coords to 3D arrays
    Grid_f                  =   (x_f, z_f);     # note
    Y_f = [];
    # Define function on coarse grid
    T                       =   cos.(pi.*X).*sin.(2*pi.*Z)

    Data_coarse1            =   tuple(T);         # 1 dataset on grid
    Data_coarse2            =   (T, T);           # 2 datasets on grid

    Grid_fine               =   (X_f, Z_f)

    Data_fine1              =   (X_f*0,);
    Data_fine2              =   (X_f.*0, Z_f.*0);

  elseif Dimension=="3D"
      # Model parameters
      W,L,H                 =   1., 1., 1.;                                    # Width, Length, Height

      # Define coarse grid
      Nx, Ny, Nz              =   33,33,33;                                                    # resolution of coarse grid

      Nx, Ny, Nz              =   150,150,150;                                                  # resolution of coarse grid

      dx,dy,dz                =   W/(Nx-1), L/(Ny-1), H/(Nz-1);                                 # grid size [m]
      x,y,z                   =   0:dx:((Nx-1)*dx),  0:dy:((Ny-1)*dy), 0:dz:((Nz-1)*dz);        # 1D coordinate arrays
      coords                  =   collect(Iterators.product(x,y,z))                             # generate coordinates from 1D coordinate vectors
      X,Y,Z                   =   (x->x[1]).(coords), (x->x[2]).(coords), (x->x[3]).(coords);   # transfer coords to 3D arrays
      Grid, Spacing           =   (x,y,z), (dx,dy,dz);

      # Define fine grid
      Nx_f, Ny_f, Nz_f        =   65, 65, 65;                                                       # resolution of fine grid
      Nx_f, Ny_f, Nz_f        =   165, 165, 165;                                                       # resolution of fine grid
      dx_f,dy_f,dz_f          =   W/(Nx_f-1), L/(Ny_f-1), H/(Nz_f-1);                               # grid size [m]
      x_f,y_f,z_f             =   0:dx_f:H, 0:dy_f:L,0:dz_f:H;                                      # 1D coordinate arrays
      coords_f                =   collect(Iterators.product(x_f,y_f, z_f))                          # generate coordinates from 1D coordinate vectors

      X_f,Y_f, Z_f            =   (x->x[1]).(coords_f), (x->x[2]).(coords_f), (x->x[3]).(coords_f); # transfer coords to 3D arrays
      Grid_f                  =   (x_f, y_f, z_f);

      # Define function on coarse grid
      T                       =   cos.(pi.*X).*sin.(2*pi.*Z).*sin.(2*pi.*Y)

      Data_coarse1            =   tuple(T);         # 1 dataset on grid
      Data_coarse2            =   (T, T, T);        # 3 datasets on grid

      Grid_fine               =   (X_f, Y_f, Z_f)
      Data_fine1              =   (X_f.*0,);
      Data_fine2              =   (X_f.*0, Y_f.*0, Z_f.*0);

  end

  # Perform interpolation from coarse-fine grid
  Interpolate!(Data_fine1, Grid, Data_coarse1, Grid_fine,   InterpolationMethod);       # one field
  Interpolate!(Data_fine2, Grid, Data_coarse2, Grid_fine,   InterpolationMethod);       # several fields

  # Compute error
  if Dimension=="2D"
    Tanal       =   cos.(pi.*X_f).*sin.(2*pi.*Z_f);
    Terror      =   Tanal - Data_fine2[2];

    if CreatePlots
      p1          =   contourf(x, z,      Data_coarse1[1]',       aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Coarse ",  dpi=150, levels=10)
      p2          =   contourf(x_f, z_f,  Data_fine2[2]',         aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Fine ",    dpi=150, levels=10)
      p3          =   contourf(x_f, z_f,  Terror',                aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="error ",    dpi=150, levels=10)

      plot(p1,p2,p3);

      png("Interpolation_2D_$InterpolationMethod")
    end

  elseif Dimension=="3D"
    Tanal       =   cos.(pi.*X_f).*sin.(2*pi.*Z_f).*sin.(2*pi.*Y_f);
    Terror      =   Tanal - Data_fine2[2];

    if CreatePlots
      p1          =   contourf(x, z,      Data_coarse1[1][:,10,:]', aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Coarse ",  dpi=150, levels=10)
      p2          =   contourf(x_f, z_f,  Data_fine2[2][:,10,:]',   aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Fine ",    dpi=150, levels=10)
      p3          =   contourf(x_f, z_f,  Terror[:,10,:]',          aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="error ",    dpi=150, levels=10)

      plot(p1,p2,p3);
      png("Interpolation_3D_$InterpolationMethod")
    end

  end

  error = norm(Terror[:],2)/length(Terror[:]);
  return error;        # return error
end

function test_SemiLagrangian2D(Method="ConstantZ",  InterpolationMethod="Linear", AdvectionMethod="RK2")
  # SemiLagrangian advection test in 2D

  # Model parameters
  W,H                     =   1.,  1.;                      # Width, Length, Height
  Nx, Nz                  =   129, 129;                      # resolution
  σ                       =   0.1;                          # halfwidth of gaussian
  Tmax                    =   2;                            # max. of gaussian
  dx,dz                   =   W/(Nx-1), H/(Nz-1);           # grid size [m]

  # Array initializations (1 - main arrays on which we can initialize properties)
  T                       =   @zeros(Nx,Nz);
  Tnew                    =   @zeros(Nx,Nz);

  # Set up model geometry & initial T structure
  x,z                     =   0:dx:((Nx-1)*dx), 0:dz:((Nz-1)*dz);
  coords                  =   collect(Iterators.product(x,z))             # generate coordinates from 1D coordinate vectors
  X,Z                     =   (x->x[1]).(coords), (x->x[2]).(coords);     # transfer coords to 2D arrays
  Grid,GridFull,Spacing   =   (x,z), (X,Z), (dx,dz);

  Xc                      =   0.5;
  Zc                      =   0.75;
  if      Method=="ConstantZ"
    # Simple advection in z-direction
    Vx        =   X.*0;
    Vz        =   Z.*0 .- 1;
    TotalTime =   0.5;                         # thermal cooling age

  elseif  Method=="Rotation"
    # Rigid body rotation:
    Vx        =    (Z .- 0.5);
    Vz        =   -(X .- 0.5);
    TotalTime =   2*pi;

  elseif  Method=="Shear"
    # Shear velocity field:
    Vx        =  -sin.(pi.*X).*cos.(pi.*Z);
    Vz        =   cos.(pi.*X).*sin.(pi.*Z);

  end

  T                      .=   Tmax.*exp.(  -( (X .- Xc).^2 .+  (Z .- Zc).^2 )./σ^2);                 # initial gaussian profile
  Tnew                   .=   T;

  maxVel      =   max( maximum(abs.(Vx)), maximum(abs.(Vz)) );
  dt          =   min(dx,dz)/2.0/maxVel;   # stable timestep (to ensure we move no more than 0.5*(dx|dz) per timesstep )
  numTime     =   ceil(TotalTime/dt);
  nt          =   Int(numTime);
  dt          =   TotalTime/nt;

  #ENV["GKSwstype"]="nul"; if isdir("viz2D_out")==false mkdir("viz2D_out") end; loadpath = "./viz2D_out/"; anim = Animation(loadpath,String[])
  #println("Animation directory: $(anim.dir)")

  time,time_kyrs          = 0.0, 0.0;
  err = 100;
  it = 0;
  #nt = 100
  for it=1:nt

      # Perform an advection step for temperature
      Tnew = AdvectTemperature(T,    Grid,  GridFull, (Vx,Vz),    dt, AdvectionMethod, InterpolationMethod);

      T, Tnew     =   Tnew, T;                                                                # Update temperature
      time        =   time + dt;                                            # Keep track of evolved time

      if mod(it,1000)==0  # print progress
          println(" Timestep $it = $((time)) ")
          #p1          =   heatmap(x_km, z_km, T[:,Int(Ny/2),:]',         aspect_ratio=1, xlims=(x_km[1],x_km[end]), ylims=(z_km[1],z_km[end]),   c=:inferno, title="Temperature, $(round(time_kyrs, digits=2)) kyrs",  dpi=150)
          p1          =   contourf(x, z, T[:,:]',         aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Temperature, $(round(time, digits=2)) ",  dpi=150, levels=10)

          plot(p1); frame(anim)
      end

  end

  x_km, z_km  =   x, z;

  # compute analytical solution
  if      Method=="ConstantZ"
    # Simple advection in z-direction
    Xc    =   0.5;
    Zc    =   0.25;

  elseif      Method=="Rotation"
      # Should arrive @ same point
  end
  Tanal =   Tmax.*exp.(  -( (X .- Xc).^2 .+  (Z .- Zc).^2 )./σ^2);                 # initial gaussian profile
  fname       = "Advection_2D_$(Method)_$(AdvectionMethod)_$(InterpolationMethod)";

  Terror      =  T - Tanal;   # error

  if CreatePlots
    # create plot
    p1          =   contourf(x, z, T',      aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Tnumeric, $(round(time, digits=2)) ",  dpi=300, levels=10)
    p2          =   contourf(x, z, Tanal',  aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Tanal, $(round(time, digits=2)) ",     dpi=300, levels=10)
    p3          =   contourf(x, z, Terror', aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Terror, $(round(time, digits=2)) ",    dpi=300, levels=10)
    plot(p1,p2,p3);
    png(fname)
  end

  error = norm(Terror[:],2)/length(Terror[:]);
  return error;        # return error
end

function test_SemiLagrangian3D(Method="ConstantZ",  InterpolationMethod="Linear", AdvectionMethod="RK2")
    # SemiLagrangian advection test in 3D

    # Model parameters
    W,L,H                   =   1., 1., 1.;                         # Width, Length, Height
    Nx, Ny, Nz              =   50, 50, 50;                         # resolution
    σ                       =   0.1;                                # halfwidth of gaussian
    Tmax                    =   2;                                  # max. of gaussian
    TotalTime               =   1;                                  # thermal cooling age
    dx,dy,dz                =   W/(Nx-1), L/(Ny-1), H/(Nz-1);       # grid size [m]
    dt                      =   min(dx,dy,dz)./2;                   # stable timestep (required for explicit FD)

    numTime                 =   ceil(TotalTime/dt);
    nt                      =   Int(numTime);

    # Array initializations (1 - main arrays on which we can initialize properties)
    T                       =   @zeros(Nx,Ny,Nz);
    Tnew                    =   @zeros(Nx,Ny,Nz);

    # Set up model geometry & initial T structure
    x,y,z                   =   0:dx:((Nx-1)*dx), 0:dy:((Ny-1)*dy), 0:dz:((Nz-1)*dz);
    coords                  =   collect(Iterators.product(x,y,z))                               # generate coordinates from 1D coordinate vectors
    X,Y,Z                   =   (x->x[1]).(coords), (x->x[2]).(coords), (x->x[3]).(coords);     # transfer coords to 3D arrays
    Grid, GridFull, Spacing =   (x,y,z), (X,Y,Z), (dx,dy,dz);
    Vx                      =  -sin.(pi.*X).*cos.(pi.*Z);
    Vy                      =   zeros(size(X));
    Vz                      =   cos.(pi.*X).*sin.(pi.*Z);

    Xc                      =   0.5;
    Yc                      =   0.5;
    Zc                      =   0.75;
    if      Method=="ConstantZ"
      # Simple advection in z-direction
      Vx        =   X.*0;
      Vy        =   Y.*0;
      Vz        =   Z.*0 .- 1;
      TotalTime =   0.5;                         # thermal cooling age

    elseif  Method=="Rotation_alongY"
      # Rigid body rotation:
      Vx        =    (Z .- 0.5);
      Vy        =   X.*0;
      Vz        =   -(X .- 0.5);
      TotalTime =   2*pi;

    elseif  Method=="Rotation_alongX"
      # Rigid body rotation:
      Vy        =    (Z .- 0.5);
      Vx        =   Y.*0;
      Vz        =   -(Y .- 0.5);
      TotalTime =   2*pi;

    elseif  Method=="Rotation_alongZ"
      # Rigid body rotation:
      Vy        =   -(X .- 0.5);
      Vx        =    (Y .- 0.5);
      Vz        =   Z.*0.0;
      TotalTime =   2*pi;
      Xc                      =   0.5;
      Yc                      =   0.75;
      Zc                      =   0.5;
   end

    T           .=   Tmax.*exp.(  -(((X .- Xc).^2 .+ (Y .- Yc).^2 .+ (Z .- Zc).^2)./(σ^2)) );                 # initial gaussian profile
    Tnew        .=   T;

    maxVel      =   max( maximum(abs.(Vx)), maximum(abs.(Vy)), maximum(abs.(Vz)) );
    dt          =   min(dx,dy,dz)/2.0/maxVel;   # stable timestep (to ensure we move no more than 0.5*(dx|dz) per timesstep )
    numTime     =   ceil(TotalTime/dt);
    nt          =   Int(numTime);
    dt          =   TotalTime/nt;

    #ENV["GKSwstype"]="nul"; if isdir("viz2D_out")==false mkdir("viz2D_out") end; loadpath = "./viz2D_out/"; anim = Animation(loadpath,String[])
    #println("Animation directory: $(anim.dir)")

    time          = 0.0;
    for it=1:nt

        # Perform an advection step for temperature

        Tnew = AdvectTemperature(T,    Grid,  GridFull, (Vx,Vy,Vz),    dt, AdvectionMethod, InterpolationMethod);

        T, Tnew         =   Tnew, T;                                                                # Update temperature
        time            =   time + dt;                                            # Keep track of evolved time

        if mod(it,1000)==0  # print progress
            println(" Timestep $it = $((time)) ")
            #p1          =   heatmap(x_km, z_km, T[:,Int(Ny/2),:]',         aspect_ratio=1, xlims=(x_km[1],x_km[end]), ylims=(z_km[1],z_km[end]),   c=:inferno, title="Temperature, $(round(time_kyrs, digits=2)) kyrs",  dpi=150)
            #p1          =   heatmap(x, z, T[:,Int(Ny/2),:]',         aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Temperature, $time",  dpi=150)
         #   p1          =   heatmap(x, y, T[:,:,Int(Nz/2)]',         aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Temperature, $time",  dpi=150)

           # p1          =   heatmap(x, z, T[Int(Nx/2),:,:]',         aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Temperature, $time",  dpi=150)

          #  plot(p1); frame(anim)
        end

    end

    x_km, z_km  =   x, z;

    # compute analytical solution
    if      Method=="ConstantZ"
      # Simple advection in z-direction
      Xc    =   0.5;
      Yc    =   0.5;
      Zc    =   0.25;

    elseif  (Method=="Rotation_alongY") || (Method=="Rotation_alongX")
      # Should arrive @ same point

    end
    Tanal =   Tmax.*exp.(  -( (X .- Xc).^2 .+  (Y .- Yc).^2 .+  (Z .- Zc).^2 )./σ^2);                 # gaussian profile
    fname = "Advection_3D_$(Method)_$(AdvectionMethod)_$(InterpolationMethod)";

    Terror      =  T - Tanal;   # error

    Tslice  = T[:,Int(Ny/2),:];
    Tanal1  = Tanal[:,Int(Ny/2),:];
    Terror1 = Terror[:,Int(Ny/2),:];

    Vx1 = Vx[:,Int(Ny/2),:];
    Vy1 = Vy[:,Int(Ny/2),:];
    Vz1 = Vz[:,Int(Ny/2),:];
    X1  = X[:,Int(Ny/2),:];
    Y1  = Y[:,Int(Ny/2),:];
    Z1  = Z[:,Int(Ny/2),:];

    if CreatePlots
      # create plot
      p1          =   contourf(x, z, Tslice',      aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Tnumeric, $(round(time, digits=2)) ",  dpi=300, levels=10)
      p2          =   contourf(x, z, Tanal1',  aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Tanal 3D, $(round(time, digits=2)) ",     dpi=300, levels=10)
      p3          =   contourf(x, z, Terror1', aspect_ratio=1, xlims=(x[1],x[end]), ylims=(z[1],z[end]),   c=:inferno, title="Terror, $(round(time, digits=2)) ",    dpi=300, levels=10)
      plot(p1,p2,p3);

      # st=20;
      # quiver!(X1[1:st:end], Z1[1:st:end], gradient=(Vx1[1:st:end]*0.1,Vz1[1:st:end]*0.1), arrow = :arrow, color = :white)

      png(fname)
    end

    error = norm(Terror[:],2)/length(Terror[:]);
    return error;        # return error
end

function test_AdvectTracers2D(Method="ConstantZ",  AdvectionMethod="RK2")
  # Test tracer advection in 2D

  # Model parameters
  W,H                     =   1.,  1.;                      # Width, Length, Height
  Nx, Nz                  =   129, 129;                     # resolution
  dx,dz                   =   W/(Nx-1), H/(Nz-1);           # grid size [m]

  # Set up model geometry
  x,z                     =   0:dx:((Nx-1)*dx), 0:dz:((Nz-1)*dz);
  coords                  =   collect(Iterators.product(x,z))             # generate coordinates from 1D coordinate vectors
  X,Z                     =   (x->x[1]).(coords), (x->x[2]).(coords);     # transfer coords to 2D arrays
  Grid, Spacing           =   (x,z), (dx,dz);

  # generate tracer coordinates
  Nxt, Nzt                =   50, 50          # num of tracers in x,z
  Wt,Ht                   =   0.1, 0.1;
  dxt,dzt                 =   Wt/(Nxt-1), Ht/(Nzt-1);
  xt,zt                   =   -Wt/2:dxt:Wt/2, -Ht/2:dz:Ht/2;
  coordst                 =   collect(Iterators.product(xt,zt))             # generate coordinates from 1D coordinate vectors
  Xt,Zt                   =   (x->x[1]).(coordst), (x->x[2]).(coordst);     # transfer coords to 2D arrays

  Xc                      =   0.5;
  Zc                      =   0.75;
  if      Method=="ConstantZ"
    # Simple advection in z-direction
    Vx        =   X.*0;
    Vz        =   Z.*0 .- 1;
    TotalTime =   0.5;                         # thermal cooling age

  elseif  Method=="Rotation"
    # Rigid body rotation:
    Vx        =    (Z .- 0.5);
    Vz        =   -(X .- 0.5);
    TotalTime =   2*pi;

  elseif  Method=="Shear"
    # Shear velocity field:
    Vx        =  -sin.(pi.*X).*cos.(pi.*Z);
    Vz        =   cos.(pi.*X).*sin.(pi.*Z);

  end

  maxVel      =   max( maximum(abs.(Vx)), maximum(abs.(Vz)) );
  dt          =   min(dx,dz)/2.0/maxVel;   # stable timestep (to ensure we move no more than 0.5*(dx|dz) per timesstep )
  numTime     =   ceil(TotalTime/dt);
  nt          =   Int(numTime);
  dt          =   TotalTime/nt;

  # Create tracer structure
  Xt .=  Xt .+ Xc; Zt .=  Zt .+ Zc;   # shift coordinates of tracers
  Xt0, Zt0    =   Xt, Zt;

  new_tracer  =   Tracer(num=1, coord=[Xt[1]; Zt[1]], T=0.);            # Create new tracer
  Tracers     =   StructArray([new_tracer]);                            # Create tracer array
  for i=firstindex(Xt)+1:lastindex(Xt)
    new_tracer  =   Tracer(num=i, coord=[Xt[i]; Zt[i]], T=0.);          # Create new tracer
    push!(Tracers, new_tracer);                             # Add new point to existing array
  end
  Tracers0 = copy(Tracers);         # create a copy of the original tracers

  #  ENV["GKSwstype"]="nul"; if isdir("viz2D_out")==false mkdir("viz2D_out") end; loadpath = "./viz2D_out/"; anim = Animation(loadpath,String[])
  #  println("Animation directory: $(anim.dir)")

  time,time_kyrs          = 0.0, 0.0;
  err = 100;
  it = 0;
  #nt = 100
  for it=1:nt

      # Perform an advection step for temperature
      AdvectTracers!(Tracers, Grid, (Vx,Vz), dt, AdvectionMethod);

      time        =   time + dt;                                            # Keep track of evolved time

      if mod(it,1000)==0  # print progress
          println(" Timestep $it = $((time)) ")

          Tr_coord = Tracers.coord; Tr_coord = hcat(Tr_coord...)';       # extract array with coordinates of tracers
          p1 = scatter(Tr_coord[:,1], Tr_coord[:,2], zcolor = Tracers.T, m = (:inferno , 0.8, Plots.stroke(0.01, :black)), markersize=1.0, xlims=(0,1), ylims=(0,1))
          plot(p1); frame(anim)
      end

  end

  x_km, z_km  =   x, z;

  # compute analytical solution
  if      Method=="ConstantZ"
    # Analytical location of tracers
    Tr_coord_anal = Tracers0.coord; Tr_coord_anal = hcat(Tr_coord_anal...)';       # extract array with coordinates of tracers
    Tr_coord_anal[:,2] .= Tr_coord_anal[:,2] .- 0.5;

  elseif      Method=="Rotation"
      # Should arrive @ same point
      Tr_coord_anal = Tracers0.coord; Tr_coord_anal = hcat(Tr_coord_anal...)';       # extract array with coordinates of tracers

  end
  Tr_coord = Tracers.coord; Tr_coord = hcat(Tr_coord...)';

  fname       = "TracersAdvection_2D_$(Method)_$(AdvectionMethod)";

  Terror      =  Tr_coord - Tr_coord_anal;   # error

  if CreatePlots
    # create plot
    p1 = plot(Tr_coord_anal[:,1], Tr_coord_anal[:,2], seriestype = :scatter, markersize=5.0, markershape=:circle, markerstrokecolor=:red, markercolor=:white, linewidth=0, label="analytics")
        plot!(Tr_coord[:,1],      Tr_coord[:,2],      seriestype = :scatter, markersize=2.0, markershape=:circle, markerstrokecolor=:black, linewidth=0, label="numerics")
    plot(p1);
    png(fname)
  end

  error = norm(Terror[:],2)/length(Terror[:]);
  return error;        # return error
end

function test_AdvectTracers3D(Method="ConstantZ",  AdvectionMethod="RK2")
  # Test tracer advection in 3D

  # Model parameters
  W,L,H                   =   1.,  1., 1.;                      # Width, Length, Height
  Nx,Ny, Nz               =   65, 65, 65;                     # resolution
  dx,dy,dz                =   W/(Nx-1), L/(Ny-1),H/(Nz-1);           # grid size [m]

  # Set up model geometry
  x,y,z                   =   0:dx:((Nx-1)*dx), 0:dy:((Ny-1)*dy), 0:dz:((Nz-1)*dz);
  coords                  =   collect(Iterators.product(x,y,z))                               # generate coordinates from 1D coordinate vectors
  X,Y,Z                   =   (x->x[1]).(coords), (x->x[2]).(coords), (x->x[3]).(coords);     # transfer coords to 2D arrays
  Grid, Spacing           =   (x,y,z), (dx,dy,dz);

  # generate tracer coordinates
  Nxt, Nyt, Nzt           =   50, 50, 50          # num of tracers in x,z
  Wt,Lt,Ht                =   0.1, 0.1, 0.1;
  dxt,dyt,dzt             =   Wt/(Nxt-1), Lt/(Nyt-1), Ht/(Nzt-1);
  xt,yt,zt                =   -Wt/2:dxt:Wt/2, -Lt/2:dyt:Lt/2, -Ht/2:dz:Ht/2;
  coordst                 =   collect(Iterators.product(xt,yt,zt))             # generate coordinates from 1D coordinate vectors
  Xt,Yt,Zt                =   (x->x[1]).(coordst), (x->x[2]).(coordst),(x->x[3]).(coordst);     # transfer coords to 2D arrays

  Xc                      =   0.5;
  Yc                      =   0.5;
  Zc                      =   0.75;
  if      Method=="ConstantZ"
    # Simple advection in z-direction
    Vx        =   X.*0;
    Vy        =   Y.*0;
    Vz        =   Z.*0 .- 1;
    TotalTime =   0.5;                         # thermal cooling age

  elseif  Method=="Rotation_alongY"
    # Rigid body rotation:
    Vx        =    (Z .- 0.5);
    Vy        =   X.*0;
    Vz        =   -(X .- 0.5);
    TotalTime =   2*pi;

  elseif  Method=="Rotation_alongX"
    # Rigid body rotation:
    Vy        =    (Z .- 0.5);
    Vx        =   Y.*0;
    Vz        =   -(Y .- 0.5);
    TotalTime =   2*pi;

  elseif  Method=="Rotation_alongZ"
    # Rigid body rotation:
    Vy        =   -(X .- 0.5);
    Vx        =    (Y .- 0.5);
    Vz        =   Z.*0.0;
    TotalTime =   2*pi;
    Xc                      =   0.5;
    Yc                      =   0.75;
    Zc                      =   0.5;
  end

  maxVel      =   max( maximum(abs.(Vx)), maximum(abs.(Vy)), maximum(abs.(Vz)) );
  dt          =   min(dx,dy,dz)/2.0/maxVel;   # stable timestep (to ensure we move no more than 0.5*(dx|dz) per timesstep )
  numTime     =   ceil(TotalTime/dt);
  nt          =   Int(numTime);
  dt          =   TotalTime/nt;

  # Create tracer structure
  Xt .=  Xt .+ Xc; Yt .=  Yt .+ Yc; Zt .=  Zt .+ Zc;   # shift coordinates of tracers

  new_tracer  =   Tracer(num=1, coord=[Xt[1]; Yt[1]; Zt[1]], T=0.);            # Create new tracer
  Tracers     =   StructArray([new_tracer]);                # Create tracer array
  for i=firstindex(Xt)+1:lastindex(Xt)
    new_tracer  =   Tracer(num=i, coord=[Xt[i]; Yt[i]; Zt[i]], T=0.);          # Create new tracer
    push!(Tracers, new_tracer);                             # Add new point to existing array
  end
  Tracers0 = copy(Tracers);         # create a copy of the original tracers

  time    = 0.0;
  err     = 100;
  it = 0;
  #nt = 100
  for it=1:nt

      # Perform an advection step for temperature
     AdvectTracers!(Tracers, Grid, (Vx,Vy,Vz), dt, AdvectionMethod);

      time        =   time + dt;                                            # Keep track of evolved time

      if mod(it,1000)==0  # print progress
          println(" Timestep $it = $((time)) ")

          Tr_coord = Tracers.coord; Tr_coord = hcat(Tr_coord...)';       # extract array with coordinates of tracers
          p1 = scatter(Tr_coord[:,1], Tr_coord[:,3], zcolor = Tracers.T, m = (:inferno , 0.8, Plots.stroke(0.01, :black)), markersize=1.0, xlims=(0,1), ylims=(0,1))
          plot(p1);
      end

  end

  x_km, z_km  =   x, z;

  # compute analytical solution
  if      Method=="ConstantZ"
    # Analytical location of tracers
    Tr_coord_anal = Tracers0.coord; Tr_coord_anal = hcat(Tr_coord_anal...)';       # extract array with coordinates of tracers
    Tr_coord_anal[:,3] .= Tr_coord_anal[:,3] .- 0.5;

  else

      # Should arrive @ same point
      Tr_coord_anal = Tracers0.coord; Tr_coord_anal = hcat(Tr_coord_anal...)';       # extract array with coordinates of tracers

  end
  Tr_coord = Tracers.coord; Tr_coord = hcat(Tr_coord...)';

  fname       = "TracersAdvection_3D_$(Method)_$(AdvectionMethod)";

  Terror      =  Tr_coord - Tr_coord_anal;   # error

  if CreatePlots
    # create plot
    p1 = plot(Tr_coord_anal[:,1], Tr_coord_anal[:,3], seriestype = :scatter, markersize=5.0, markershape=:circle, markerstrokecolor=:red, markercolor=:white, linewidth=0, label="analytics")
        plot!(Tr_coord[:,1],      Tr_coord[:,3],      seriestype = :scatter, markersize=2.0, markershape=:circle, markerstrokecolor=:black, linewidth=0, label="numerics")
    plot(p1);
    png(fname)
  end

  error = norm(Terror[:],2)/length(Terror[:]);
  return error;        # return error
end

# ===================================================================================================
if 1==1

@testset "Interpolation" begin
  @test test_Interpolation("2D", "Linear")      ≈  2.9880072526933544e-5    atol=1e-8;
  @test test_Interpolation("2D", "Quadratic")   ≈  9.72101621785853e-7      atol=1e-8;
  @test test_Interpolation("2D", "Cubic")       ≈  5.752004968381702e-7     atol=1e-8;
  @test test_Interpolation("3D", "Linear")      ≈  5.758940704356386e-8     atol=1e-8;
  @test test_Interpolation("3D", "Cubic")       ≈  1.0924731223751025e-10   atol=1e-8;
end;

@testset "2D semi-lagrangian advection" begin
  @test test_SemiLagrangian2D("ConstantZ","Cubic", "RK4"    )   ≈ 7.257090264260951e-6    atol=1e-8;
  @test test_SemiLagrangian2D("ConstantZ","Linear","Euler"  )   ≈ 0.00026385729240145655  atol=1e-6;

  @test test_SemiLagrangian2D("Rotation", "Linear", "Euler" )   ≈ 0.0009061021221256843   atol=1e-8;
  @test test_SemiLagrangian2D("Rotation", "Linear", "RK2"   )   ≈ 0.000892833826724465    atol=1e-8;
  @test test_SemiLagrangian2D("Rotation", "Linear", "RK4"   )   ≈ 0.0008928336158448796   atol=1e-8;

  @test test_SemiLagrangian2D("Rotation", "Quadratic", "RK2")   ≈ 2.2477968913260075e-5   atol=1e-8;

  @test test_SemiLagrangian2D("Rotation", "Cubic", "Euler"  )   ≈ 0.00013400741577631413  atol=1e-8;
  @test test_SemiLagrangian2D("Rotation", "Cubic", "RK2"    )   ≈ 1.0762102257417536e-6   atol=1e-8;
  @test test_SemiLagrangian2D("Rotation", "Cubic", "RK4"    )   ≈ 1.069882664228221e-6    atol=1e-8;
end;

@testset "3D semi-lagrangian advection" begin
  @test test_SemiLagrangian3D("ConstantZ","Linear","RK2"   )          ≈ 6.822115778040405e-5    atol=1e-8;
  @test test_SemiLagrangian3D("ConstantZ","Cubic", "RK2"   )          ≈ 1.1038374380141209e-6   atol=1e-8;

  @test test_SemiLagrangian3D("Rotation_alongY","Linear", "Euler")    ≈ 0.00017064553949646487  atol=1e-8;
  @test test_SemiLagrangian3D("Rotation_alongY","Cubic", "RK2"   )    ≈ 1.899516467222456e-6    atol=1e-8;
  @test test_SemiLagrangian3D("Rotation_alongY","Cubic", "RK2"   )    ≈ 1.899516467222456e-6    atol=1e-8;

  @test test_SemiLagrangian3D("Rotation_alongX","Cubic", "RK2"   )    ≈ 1.8995164672224572e-6    atol=1e-8;
  @test test_SemiLagrangian3D("Rotation_alongY","Cubic", "RK2"   )    ≈ 1.899516467222456e-6    atol=1e-8;
  @test test_SemiLagrangian3D("Rotation_alongZ","Cubic", "RK2"   )    ≈ 1.8995164672224786e-6   atol=1e-8;
end;

@testset "2D tracer advection" begin
  @test test_AdvectTracers2D("Rotation","Euler")        ≈ 0.00012187143210808163  atol=1e-8;
  @test test_AdvectTracers2D("Rotation","RK2")          ≈ 3.132104548068979e-7    atol=1e-8;
  @test test_AdvectTracers2D("Rotation","RK4")          ≈ 1.5661110135721713e-7   atol=1e-8;

  @test test_AdvectTracers2D("ConstantZ","Euler")       ≈ 0.0   atol=1e-8;
end;

@testset "3D tracer advection" begin
  @test test_AdvectTracers3D("ConstantZ","Euler")       ≈ 0.0   atol=1e-8;

  @test test_AdvectTracers3D("Rotation_alongX","RK2")   ≈ 1.60725931984687e-7    atol=1e-8;
  @test test_AdvectTracers3D("Rotation_alongY","RK2")   ≈ 1.607259319846869e-7   atol=1e-8;
  @test test_AdvectTracers3D("Rotation_alongZ","RK2")   ≈ 1.6256291449684137e-7  atol=1e-8;

  @test test_AdvectTracers3D("Rotation_alongY","Euler") ≈ 3.169439899086836e-5   atol=1e-8;
  @test test_AdvectTracers3D("Rotation_alongY","RK4")   ≈ 8.037493317907375e-8   atol=1e-8;
end;

end

#@time test_SemiLagrangian3D("ConstantZ","Cubic", "RK4"    )
#test_SemiLagrangian3D("ConstantZ","Linear", "Euler"    )
#test_SemiLagrangian2D("ConstantZ","Linear", "Euler"    )
#test_AdvectTracers2D("Rotation","Euler")
#test_AdvectTracers2D("ConstantZ","Euler")

#test_AdvectTracers3D("Rotation_alongX","RK2")

#test_SemiLagrangian2D("Rotation", "Linear", "RK2"   )


using ModelingToolkit
@variables x y
@parameters a b
loss = (a - x)^2 + b * (y - x^2)^2
sys1 = OptimizationSystem(loss,[x,y],[a,b],name=:sys1)
sys2 = OptimizationSystem(loss,[x,y],[a,b],name=:sys2)

@variables z
@parameters β
loss2 = sys1.x - sys2.y + z*β
combinedsys = OptimizationSystem(loss2,[z],[β],systems=[sys1,sys2],name=:combinedsys)

equations(combinedsys)
states(combinedsys)
parameters(combinedsys)

calculate_gradient(combinedsys)
calculate_hessian(combinedsys)
generate_function(combinedsys)
generate_gradient(combinedsys)
generate_hessian(combinedsys)


using Eigenbroetler
using Test

random(rows::Integer, cols::Integer) =
    Eigenbrot(rand(ComplexF64, rows, cols), rand(Bool))

function sample(x, y)
    θ = angle(x + y * im)
    r = hypot(x, y)
    x^2 - y^2 - im*(r * cos(θ))
end

@testset "Eigenbroetler" begin
    @testset "Constructors" begin
        eb = Eigenbrot(512, 512, rand(Bool)) do x, y
            if x == 0 && y == 0
                return 0.0
            end
            r = sqrt(x * x + y * y)
            ϕ = atan(y, x)
            return cos(30 * log(r) + 13 * ϕ) ^ 2 + cos(20 * log(r) - 11 * ϕ) ^ 2
        end
        @test height(eb) == 512
        @test width(eb) == 512
        @test size(eb, 1) == 512
        @test size(eb, 2) == 512
        @test length(eb) == 512 * 512
        @test ndims(eb) == 2
        @test size(eb) == (512, 512)
    end

    @testset "Files" begin
        eb = Eigenbrot(sample, 64, 64)
        testfile = tempname() * ".png"
        save(testfile, eb, ImageSetting(RealPart, LinearScale))
        pngmagic = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]
        t = open(testfile)
        by = read(t, length(pngmagic))
        close(t)
        @test by[1:length(pngmagic)] == pngmagic
    end

    @testset "Basics" begin
        eb = random(256, 512)
        @test height(eb) == 256
        @test width(eb) == 512
        @test size(eb, 1) == 256
        @test size(eb, 2) == 512
        @test length(eb) == 256 * 512
        @test ndims(eb) == 2
        @test size(eb) == (256, 512)
        testfile = tempname()
        save(testfile, eb)
        @test ispath(testfile)
        ebr = Eigenbrot(testfile)
        @test eb.vals ≈ ebr.vals
        @test eb.fft == ebr.fft
        rm(testfile, force = false)
    end

    @testset "Indexing" begin
        eb = Eigenbrot(50, 50)
        for c in 1:50
            for r in 1:50
                val = rand(ComplexF64)
                eb[r, c] = val
                @test val == eb[r, c]
            end
        end
        for c in 1:10
            for r in 1:10
                xy = x, y = coords(eb, r, c)
                @test pixel(eb, x, y) == (r, c)
                @test pixel(eb, xy) == (r, c)
                @test xy == coords(eb, (r, c))
            end
        end
    end

    @testset "Simple methods" begin
        eb = random(8, 8)
        ans = real(eb)
        @test real(eb.vals) ≈ ans.vals
        @test eb.fft == ans.fft
        ans = imag(eb)
        @test imag(eb.vals) ≈ ans.vals
        @test eb.fft == ans.fft
        ans = angle(eb)
        @test angle.(eb.vals) ≈ ans.vals
        @test eb.fft == ans.fft
        ans = abs(eb)
        @test abs.(eb.vals) ≈ ans.vals
        @test eb.fft == ans.fft
        ans = conj(eb)
        @test conj(eb.vals) ≈ ans.vals
        @test eb.fft == ans.fft
    end

    @testset "Arithmetic" begin
        a = random(8, 8)
        b = random(8, 8)
        ans = 1 + a
        @test ans.vals ≈ 1 .+ a.vals
        @test ans.fft == a.fft
        ans = a + 1
        @test ans.vals ≈ 1 .+ a.vals
        @test ans.fft == a.fft

        ans = 1 - a
        @test ans.vals ≈ 1 .- a.vals
        @test ans.fft == a.fft
        ans = a - 1
        @test ans.vals ≈ a.vals .- 1
        @test ans.fft == a.fft

        ans = 2 * a
        @test ans.vals ≈ 2 * a.vals
        @test ans.fft == a.fft
        ans = a * 2
        @test ans.vals ≈ 2 * a.vals
        @test ans.fft == a.fft
        ans = a / 2
        @test ans.vals ≈ a.vals / 2
        @test ans.fft == a.fft

        ans = a * b
        @test ans.vals ≈ a.vals .* b.vals
        @test ans.fft == a.fft
        ans = b * a
        @test ans.vals ≈ a.vals .* b.vals
        @test ans.fft == b.fft

        ans = a / b
        @test ans.vals ≈ a.vals ./ b.vals
        @test ans.fft == a.fft
        ans = b / a
        @test ans.vals ≈ b.vals ./ a.vals
        @test ans.fft == b.fft
    end

    @testset "Sample with padding" begin
        eb = Eigenbrot(sample, 64, 64)
        ebf = Eigenbrot(joinpath(@__DIR__, "sample.fit"))
        @test eb.vals ≈ ebf.vals
        @test eb.fft == ebf.fft

        l, t, r, b = 5, 8, 3, 7
        padvalue = rand(ComplexF64)
        padded = pad(eb, padvalue, left = l, top = t, right = r, bottom = b)
        @test size(padded) == size(eb) .+ (t + b, l + r)
        @test all(padded.vals[1:b, :] .== padvalue)
        @test all(padded.vals[end-t+1:end, :] .== padvalue)
        @test all(padded.vals[:, 1:l] .== padvalue)
        @test all(padded.vals[:, end-r+1:end] .== padvalue)
        @test padded.vals[b+1:end-t, l+1:end-r] == eb.vals
    end

    @testset "Fourier" begin
        eb = random(50, 50)
        for ft in [fftx, ffty, fft]
            ebfft = ft(eb)
            ebfft2 = ft(ebfft)
            @test eb.vals ≈ ebfft2.vals
            @test eb.fft == ebfft2.fft
        end
        nodc = removeDC(eb)
        nodcff = fft(nodc)
        @test isapprox(nodcff[pixel(nodcff, 0, 0)], 0, atol = 1e-15)

        xscale = 3//2
        yscale = 5//4
        eb = random(8, 8)

        ebs = scale_chirpz(eb, xscale, yscale)
        @test size(ebs) == (10, 12)
        ebs = scale_x_chirpz(eb, xscale)
        @test size(ebs) == (8, 12)
        ebs = scale_y_chirpz(eb, yscale)
        @test size(ebs) == (10, 8)
    end
end