text:
import Pkg; Pkg.activate(@__DIR__); Pkg.instantiate()

using BenchmarkTools
using BSON
using LoopVectorization
using MuladdMacro
using StaticArrays
using Tullio

###################################################################################################
# 2D versions
function multiply_dimensionwise_sequential!(
    data_out::AbstractArray{<:Any, 3}, data_in::AbstractArray{<:Any, 3}, vandermonde,
    tmp1=zeros(eltype(data_out), size(data_out, 1), size(vandermonde, 1), size(vandermonde, 2)))
  n_vars      = size(data_out, 1)
  n_nodes_out = size(vandermonde, 1)
  n_nodes_in  = size(vandermonde, 2)
  data_out .= zero(eltype(data_out))

  @boundscheck begin
    inbounds = (size(data_out, 2) == size(data_out, 3) == n_nodes_out) &&
               (size(data_in,  1) == n_vars) &&
               (size(data_in,  2) == size(data_in,  3) == n_nodes_in)
    inbounds || throw(BoundsError())
  end

  # Interpolate in x-direction
  @inbounds for j in 1:n_nodes_in, i in 1:n_nodes_out
    for ii in 1:n_nodes_in
      for v in 1:n_vars
        tmp1[v, i, j] += vandermonde[i, ii] * data_in[v, ii, j]
      end
    end
  end

  # Interpolate in y-direction
  @inbounds for j in 1:n_nodes_out, i in 1:n_nodes_out
    for jj in 1:n_nodes_in
      for v in 1:n_vars
        data_out[v, i, j] += vandermonde[j, jj] * tmp1[v, i, jj]
      end
    end
  end

  return data_out
end

function multiply_dimensionwise_sequential_avx!(
    data_out::AbstractArray{<:Any, 3}, data_in::AbstractArray{<:Any, 3}, vandermonde,
    tmp1=zeros(eltype(data_out), size(data_out, 1), size(vandermonde, 1), size(vandermonde, 2)))
  n_vars      = size(data_out, 1)
  n_nodes_out = size(vandermonde, 1)
  n_nodes_in  = size(vandermonde, 2)
  data_out .= zero(eltype(data_out))

  @boundscheck begin
    inbounds = (size(data_out, 2) == size(data_out, 3) == n_nodes_out) &&
               (size(data_in,  1) == n_vars) &&
               (size(data_in,  2) == size(data_in,  3) == n_nodes_in)
    inbounds || throw(BoundsError())
  end

  # Interpolate in x-direction
  @avx for j in 1:n_nodes_in, i in 1:n_nodes_out
    for ii in 1:n_nodes_in
      for v in 1:n_vars
        tmp1[v, i, j] += vandermonde[i, ii] * data_in[v, ii, j]
      end
    end
  end

  # Interpolate in y-direction
  @avx for j in 1:n_nodes_out, i in 1:n_nodes_out
    for jj in 1:n_nodes_in
      for v in 1:n_vars
        data_out[v, i, j] += vandermonde[j, jj] * tmp1[v, i, jj]
      end
    end
  end

  return data_out
end

function multiply_dimensionwise_sequential_tullio!(
    data_out::AbstractArray{<:Any, 3}, data_in::AbstractArray{<:Any, 3}, vandermonde,
    tmp1=zeros(eltype(data_out), size(data_out, 1), size(vandermonde, 1), size(vandermonde, 2)))

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

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

  return data_out
end

@generated function multiply_dimensionwise_sequential_nexpr!(
    data_out::AbstractArray{<:Any, 3}, data_in::AbstractArray{<:Any, 3}, vandermonde::SMatrix{n_nodes_out,n_nodes_in}, ::Val{n_vars}) where {n_nodes_out, n_nodes_in, n_vars}
  quote
    @boundscheck begin
      inbounds = (size(data_out, 1) == $n_vars) &&
                 (size(data_out, 2) == size(data_out, 3) == $n_nodes_out) &&
                 (size(data_in,  1) == $n_vars) &&
                 (size(data_in,  2) == size(data_in,  3) == $n_nodes_in)
      inbounds || throw(BoundsError())
    end

    # Interpolate in x-direction
    @inbounds @muladd Base.Cartesian.@nexprs $n_nodes_in j -> begin
      Base.Cartesian.@nexprs $n_nodes_out i -> begin
        Base.Cartesian.@nexprs $n_vars v -> begin
          tmp1_v_i_j = zero(eltype(data_out))
          Base.Cartesian.@nexprs $n_nodes_in ii -> begin
            tmp1_v_i_j += vandermonde[i, ii] * data_in[v, ii, j]
          end
        end
      end
    end

    # Interpolate in y-direction
    @inbounds @muladd Base.Cartesian.@nexprs $n_nodes_out j -> begin
      Base.Cartesian.@nexprs $n_nodes_out i -> begin
        Base.Cartesian.@nexprs $n_vars v -> begin
          tmp2_v_i_j = zero(eltype(data_out))
          Base.Cartesian.@nexprs $n_nodes_in jj -> begin
            tmp2_v_i_j += vandermonde[j, jj] * tmp1_v_i_jj
          end
          data_out[v, i, j] = tmp2_v_i_j
        end
      end
    end

    return data_out
  end
end

function multiply_dimensionwise_squeezed!(
    data_out::AbstractArray{<:Any, 3}, data_in::AbstractArray{<:Any, 3}, vandermonde)
  n_vars      = size(data_out, 1)
  n_nodes_out = size(vandermonde, 1)
  n_nodes_in  = size(vandermonde, 2)
  data_out .= zero(eltype(data_out))

  @boundscheck begin
    inbounds = (size(data_out, 2) == size(data_out, 3) == n_nodes_out) &&
               (size(data_in,  1) == n_vars) &&
               (size(data_in,  2) == size(data_in,  3) == n_nodes_in)
    inbounds || throw(BoundsError())
  end

  @inbounds for j in 1:n_nodes_out, i in 1:n_nodes_out
    for v in 1:n_vars
      acc = zero(eltype(data_out))
      for jj in 1:n_nodes_in, ii in 1:n_nodes_in
        acc += vandermonde[i, ii]*  vandermonde[j, jj] * data_in[v, ii, jj]
      end
      data_out[v, i, j] = acc
    end
  end

  return data_out
end

function multiply_dimensionwise_squeezed_avx!(
    data_out::AbstractArray{<:Any, 3}, data_in::AbstractArray{<:Any, 3}, vandermonde)
  n_vars      = size(data_out, 1)
  n_nodes_out = size(vandermonde, 1)
  n_nodes_in  = size(vandermonde, 2)
  data_out .= zero(eltype(data_out))

  @boundscheck begin
    inbounds = (size(data_out, 2) == size(data_out, 3) == n_nodes_out) &&
               (size(data_in,  1) == n_vars) &&
               (size(data_in,  2) == size(data_in,  3) == n_nodes_in)
    inbounds || throw(BoundsError())
  end

  @avx for j in 1:n_nodes_out, i in 1:n_nodes_out
    for v in 1:n_vars
      acc = zero(eltype(data_out))
      for jj in 1:n_nodes_in, ii in 1:n_nodes_in
        acc += vandermonde[i, ii] * vandermonde[j, jj] * data_in[v, ii, jj]
      end
      data_out[v, i, j] = acc
    end
  end

  return data_out
end

function multiply_dimensionwise_squeezed_tullio!(
    data_out::AbstractArray{<:Any, 3}, data_in::AbstractArray{<:Any, 3}, vandermonde)

  @tullio threads=false data_out[v, i, j] = vandermonde[i, ii] * vandermonde[j, jj] * data_in[v, ii, jj]

  return data_out
end

function run_benchmarks_2d(n_vars=4, n_nodes_in=4, n_nodes_out=2*n_nodes_in)
  data_in  = randn(n_vars, n_nodes_in,  n_nodes_in)
  data_out = randn(n_vars, n_nodes_out, n_nodes_out)
  vandermonde_dynamic = randn(n_nodes_out, n_nodes_in)
  vandermonde_static  = SMatrix{n_nodes_out, n_nodes_in}(vandermonde_dynamic)
  vandermonde_mmatrix = MMatrix{n_nodes_out, n_nodes_in}(vandermonde_dynamic)

  println("\n\n# 2D   ", "#"^70)
  println("n_vars      = ", n_vars)
  println("n_nodes_in  = ", n_nodes_in)
  println("n_nodes_out = ", n_nodes_out)
  println()

  println("\n","multiply_dimensionwise_sequential!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_sequential!(data_out, data_in, vandermonde_dynamic)
  data_out_copy = copy(data_out)
  display(@benchmark multiply_dimensionwise_sequential!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_sequential!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential!($(data_out), $(data_in), $(vandermonde_static)))
  println("\n", "vandermonde_mmatrix")
  multiply_dimensionwise_sequential!(data_out, data_in, vandermonde_mmatrix)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential!($(data_out), $(data_in), $(vandermonde_mmatrix)))
  println()

  println("\n","multiply_dimensionwise_sequential_avx!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_sequential_avx!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_avx!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_sequential_avx!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_avx!($(data_out), $(data_in), $(vandermonde_static)))
  println("\n", "vandermonde_mmatrix")
  multiply_dimensionwise_sequential_avx!(data_out, data_in, vandermonde_mmatrix)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_avx!($(data_out), $(data_in), $(vandermonde_mmatrix)))
  println()

  println("\n","multiply_dimensionwise_sequential_tullio!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_sequential_tullio!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_sequential_tullio!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_static)))
  println("\n", "vandermonde_mmatrix")
  multiply_dimensionwise_sequential_tullio!(data_out, data_in, vandermonde_mmatrix)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_mmatrix)))
  println()

  println("\n","multiply_dimensionwise_sequential_nexpr!")
  println("vandermonde_static")
  multiply_dimensionwise_sequential_nexpr!(data_out, data_in, vandermonde_static, Val(n_vars))
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_nexpr!($(data_out), $(data_in), $(vandermonde_static), $(Val(n_vars))))
  println()

  println("\n","multiply_dimensionwise_squeezed!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_squeezed!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_squeezed!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed!($(data_out), $(data_in), $(vandermonde_static)))
  println("\n", "vandermonde_mmatrix")
  multiply_dimensionwise_squeezed!(data_out, data_in, vandermonde_mmatrix)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed!($(data_out), $(data_in), $(vandermonde_mmatrix)))
  println()

  println("\n","multiply_dimensionwise_squeezed_avx!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_squeezed_avx!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_avx!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_squeezed_avx!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_avx!($(data_out), $(data_in), $(vandermonde_static)))
  println("\n", "vandermonde_mmatrix")
  multiply_dimensionwise_squeezed_avx!(data_out, data_in, vandermonde_mmatrix)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_avx!($(data_out), $(data_in), $(vandermonde_mmatrix)))
  println()

  println("\n","multiply_dimensionwise_squeezed_tullio!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_squeezed_tullio!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_tullio!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_squeezed_tullio!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_tullio!($(data_out), $(data_in), $(vandermonde_static)))
  println("\n", "vandermonde_mmatrix")
  multiply_dimensionwise_squeezed_tullio!(data_out, data_in, vandermonde_mmatrix)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_tullio!($(data_out), $(data_in), $(vandermonde_mmatrix)))
  println()

  nothing
end

# TODO
run_benchmarks_2d(4, 4, 4) # n_vars, n_nodes_in, n_nodes_out

function compute_benchmarks_2d(n_vars, n_nodes_in, n_nodes_out)
  data_in  = randn(n_vars, n_nodes_in,  n_nodes_in)
  data_out = randn(n_vars, n_nodes_out, n_nodes_out)
  vandermonde_dynamic = randn(n_nodes_out, n_nodes_in)
  vandermonde_static = SMatrix{n_nodes_out, n_nodes_in}(vandermonde_dynamic)
  tmp1 = zeros(eltype(data_out), n_vars, n_nodes_out, n_nodes_in)

  println("n_vars = ", n_vars, "; n_nodes_in = ", n_nodes_in, "; n_nodes_out = ", n_nodes_out)
  sequential_dynamic = @benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_dynamic))
  sequential_static  = @benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_static))
  #FIXME sequential_nexpr   = @benchmark multiply_dimensionwise_sequential_nexpr!($(data_out), $(data_in), $(vandermonde_static), $(Val(n_vars)))
  sequential_dynamic_prealloc = @benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_dynamic), $(tmp1))
  sequential_static_prealloc  = @benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_static), $(tmp1))
  squeezed_dynamic = @benchmark multiply_dimensionwise_squeezed_tullio!($(data_out), $(data_in), $(vandermonde_dynamic))
  squeezed_static  = @benchmark multiply_dimensionwise_squeezed_tullio!($(data_out), $(data_in), $(vandermonde_static))

  return time(median(sequential_dynamic)),
         time(median(sequential_static)),
         NaN, #FIXME time(median(sequential_nexpr)),
         time(median(sequential_dynamic_prealloc)),
         time(median(sequential_static_prealloc)),
         time(median(squeezed_dynamic)),
         time(median(squeezed_static))
end

function compute_benchmarks_2d(n_vars_list, n_nodes_in_list)
  sequential_dynamic          = zeros(length(n_vars_list), length(n_nodes_in_list))
  sequential_static           = zeros(length(n_vars_list), length(n_nodes_in_list))
  sequential_nexpr            = zeros(length(n_vars_list), length(n_nodes_in_list))
  sequential_dynamic_prealloc = zeros(length(n_vars_list), length(n_nodes_in_list))
  sequential_static_prealloc  = zeros(length(n_vars_list), length(n_nodes_in_list))
  squeezed_dynamic            = zeros(length(n_vars_list), length(n_nodes_in_list))
  squeezed_static             = zeros(length(n_vars_list), length(n_nodes_in_list))

  # n_nodes_out = n_nodes_in
  # mortar
  # superset of n_vars = 1, n_nodes_out = n_nodes_in, used for blending
  for (idx_nodes, n_nodes_in) in enumerate(n_nodes_in_list)
    for (idx_variables, n_vars) in enumerate(n_vars_list)
      n_nodes_out = n_nodes_in
      sequential_dynamic[idx_variables, idx_nodes],
      sequential_static[idx_variables, idx_nodes],
      sequential_nexpr[idx_variables, idx_nodes],
      sequential_dynamic_prealloc[idx_variables, idx_nodes],
      sequential_static_prealloc[idx_variables, idx_nodes],
      squeezed_dynamic[idx_variables, idx_nodes],
      squeezed_static[idx_variables, idx_nodes] =
        compute_benchmarks_2d(n_vars, n_nodes_in, n_nodes_out)
    end
  end
  BSON.@save "2D_nVarTotal_nNodesIn.bson" n_vars_list n_nodes_in_list sequential_dynamic sequential_static sequential_nexpr sequential_dynamic_prealloc sequential_static_prealloc squeezed_dynamic squeezed_static

  # n_nodes_out = 2*n_nodes_in
  # visualization
  for (idx_nodes, n_nodes_in) in enumerate(n_nodes_in_list)
    for (idx_variables, n_vars) in enumerate(n_vars_list)
      n_nodes_out = 2 * n_nodes_in
      sequential_dynamic[idx_variables, idx_nodes],
      sequential_static[idx_variables, idx_nodes],
      sequential_nexpr[idx_variables, idx_nodes],
      sequential_dynamic_prealloc[idx_variables, idx_nodes],
      sequential_static_prealloc[idx_variables, idx_nodes],
      squeezed_dynamic[idx_variables, idx_nodes],
      squeezed_static[idx_variables, idx_nodes] =
        compute_benchmarks_2d(n_vars, n_nodes_in, n_nodes_out)
    end
  end
  BSON.@save "2D_nVarTotal_2nNodesIn.bson" n_vars_list n_nodes_in_list sequential_dynamic sequential_static sequential_nexpr sequential_dynamic_prealloc sequential_static_prealloc squeezed_dynamic squeezed_static

  return nothing
end

# TODO
compute_benchmarks_2d(1:10, 2:10)

###################################################################################################
# 3D versions

function multiply_dimensionwise_sequential!(
    data_out::AbstractArray{<:Any, 4}, data_in::AbstractArray{<:Any, 4}, vandermonde,
    tmp1=zeros(eltype(data_out), size(data_out, 1), size(vandermonde, 1), size(vandermonde, 2), size(vandermonde, 2)),
    tmp2=zeros(eltype(data_out), size(data_out, 1), size(vandermonde, 1), size(vandermonde, 1), size(vandermonde, 2)))
  n_vars      = size(data_out, 1)
  n_nodes_out = size(vandermonde, 1)
  n_nodes_in  = size(vandermonde, 2)
  data_out .= zero(eltype(data_out))

  @boundscheck begin
    inbounds = (size(data_out, 2) == size(data_out, 3) == size(data_out, 4) == n_nodes_out) &&
               (size(data_in,  1) == n_vars) &&
               (size(data_in,  2) == size(data_in,  3) == size(data_in,  4) == n_nodes_in)
    inbounds || throw(BoundsError())
  end

  # Interpolate in x-direction
  @inbounds for k in 1:n_nodes_in, j in 1:n_nodes_in, i in 1:n_nodes_out
    for ii in 1:n_nodes_in
      for v in 1:n_vars
        tmp1[v, i, j, k] += vandermonde[i, ii] * data_in[v, ii, j, k]
      end
    end
  end

  # Interpolate in y-direction
  @inbounds for k in 1:n_nodes_in, j in 1:n_nodes_out, i in 1:n_nodes_out
    for jj in 1:n_nodes_in
      for v in 1:n_vars
        tmp2[v, i, j, k] += vandermonde[j, jj] * tmp1[v, i, jj, k]
      end
    end
  end

  # Interpolate in z-direction
  @inbounds for k in 1:n_nodes_out, j in 1:n_nodes_out, i in 1:n_nodes_out
    for kk in 1:n_nodes_in
      for v in 1:n_vars
        data_out[v, i, j, k] += vandermonde[k, kk] * tmp2[v, i, j, kk]
      end
    end
  end

  return data_out
end

function multiply_dimensionwise_sequential_avx!(
    data_out::AbstractArray{<:Any, 4}, data_in::AbstractArray{<:Any, 4}, vandermonde,
    tmp1=zeros(eltype(data_out), size(data_out, 1), size(vandermonde, 1), size(vandermonde, 2), size(vandermonde, 2)),
    tmp2=zeros(eltype(data_out), size(data_out, 1), size(vandermonde, 1), size(vandermonde, 1), size(vandermonde, 2)))
  n_vars      = size(data_out, 1)
  n_nodes_out = size(vandermonde, 1)
  n_nodes_in  = size(vandermonde, 2)
  data_out .= zero(eltype(data_out))

  @boundscheck begin
    inbounds = (size(data_out, 2) == size(data_out, 3) == size(data_out, 4) == n_nodes_out) &&
               (size(data_in,  1) == n_vars) &&
               (size(data_in,  2) == size(data_in,  3) == size(data_in,  4) == n_nodes_in)
    inbounds || throw(BoundsError())
  end

  # Interpolate in x-direction
  @avx for k in 1:n_nodes_in, j in 1:n_nodes_in, i in 1:n_nodes_out
    for ii in 1:n_nodes_in
      for v in 1:n_vars
        tmp1[v, i, j, k] += vandermonde[i, ii] * data_in[v, ii, j, k]
      end
    end
  end

  # Interpolate in y-direction
  @avx for k in 1:n_nodes_in, j in 1:n_nodes_out, i in 1:n_nodes_out
    for jj in 1:n_nodes_in
      for v in 1:n_vars
        tmp2[v, i, j, k] += vandermonde[j, jj] * tmp1[v, i, jj, k]
      end
    end
  end

  # Interpolate in z-direction
  @avx for k in 1:n_nodes_out, j in 1:n_nodes_out, i in 1:n_nodes_out
    for kk in 1:n_nodes_in
      for v in 1:n_vars
        data_out[v, i, j, k] += vandermonde[k, kk] * tmp2[v, i, j, kk]
      end
    end
  end

  return data_out
end

function multiply_dimensionwise_sequential_tullio!(
    data_out::AbstractArray{<:Any, 4}, data_in::AbstractArray{<:Any, 4}, vandermonde,
    tmp1=zeros(eltype(data_out), size(data_out, 1), size(vandermonde, 1), size(vandermonde, 2), size(vandermonde, 2)),
    tmp2=zeros(eltype(data_out), size(data_out, 1), size(vandermonde, 1), size(vandermonde, 1), size(vandermonde, 2)))

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

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

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

  return data_out
end

@generated function multiply_dimensionwise_sequential_nexpr!(
    data_out::AbstractArray{<:Any, 4}, data_in::AbstractArray{<:Any, 4}, vandermonde::SMatrix{n_nodes_out,n_nodes_in}, ::Val{n_vars}) where {n_nodes_out, n_nodes_in, n_vars}
  quote
    @boundscheck begin
      inbounds = (size(data_out, 1) == $n_vars) &&
                 (size(data_out, 2) == size(data_out, 3) == size(data_out, 4) == $n_nodes_out) &&
                 (size(data_in,  1) == $n_vars) &&
                 (size(data_in,  2) == size(data_in,  3) == size(data_in,  4) == $n_nodes_in)
      inbounds || throw(BoundsError())
    end

    # Interpolate in x-direction
    @inbounds @muladd Base.Cartesian.@nexprs $n_nodes_in k -> begin
      Base.Cartesian.@nexprs $n_nodes_in j -> begin
        Base.Cartesian.@nexprs $n_nodes_out i -> begin
          Base.Cartesian.@nexprs $n_vars v -> begin
            tmp1_v_i_j_k = zero(eltype(data_out))
            Base.Cartesian.@nexprs $n_nodes_in ii -> begin
              tmp1_v_i_j_k += vandermonde[i, ii] * data_in[v, ii, j, k]
            end
          end
        end
      end
    end

    # Interpolate in y-direction
    @inbounds @muladd Base.Cartesian.@nexprs $n_nodes_in k -> begin
      Base.Cartesian.@nexprs $n_nodes_out j -> begin
        Base.Cartesian.@nexprs $n_nodes_out i -> begin
          Base.Cartesian.@nexprs $n_vars v -> begin
            tmp2_v_i_j_k = zero(eltype(data_out))
            Base.Cartesian.@nexprs $n_nodes_in jj -> begin
              tmp2_v_i_j_k += vandermonde[j, jj] * tmp1_v_i_jj_k
            end
          end
        end
      end
    end

    # Interpolate in z-direction
    @inbounds @muladd Base.Cartesian.@nexprs $n_nodes_out k -> begin
      Base.Cartesian.@nexprs $n_nodes_out j -> begin
        Base.Cartesian.@nexprs $n_nodes_out i -> begin
          Base.Cartesian.@nexprs $n_vars v -> begin
            tmp3_v_i_j_k = zero(eltype(data_out))
            Base.Cartesian.@nexprs $n_nodes_in kk -> begin
              tmp3_v_i_j_k += vandermonde[k, kk] * tmp2_v_i_j_kk
            end
            data_out[v, i, j, k] = tmp3_v_i_j_k
          end
        end
      end
    end

    return data_out
  end
end

function multiply_dimensionwise_squeezed!(
    data_out::AbstractArray{<:Any, 4}, data_in::AbstractArray{<:Any, 4}, vandermonde)
  n_vars      = size(data_out, 1)
  n_nodes_out = size(vandermonde, 1)
  n_nodes_in  = size(vandermonde, 2)
  data_out .= zero(eltype(data_out))

  @boundscheck begin
    inbounds = (size(data_out, 2) == size(data_out, 3) == size(data_out, 4) == n_nodes_out) &&
               (size(data_in,  1) == n_vars) &&
               (size(data_in,  2) == size(data_in,  3) == size(data_in,  4) == n_nodes_in)
    inbounds || throw(BoundsError())
  end

  @inbounds for k in 1:n_nodes_out, j in 1:n_nodes_out, i in 1:n_nodes_out
    for v in 1:n_vars
      acc = zero(eltype(data_out))
      for kk in 1:n_nodes_in, jj in 1:n_nodes_in, ii in 1:n_nodes_in
        acc += vandermonde[i, ii] * vandermonde[j, jj] * vandermonde[k, kk] * data_in[v, ii, jj, kk]
      end
      data_out[v, i, j, k] = acc
    end
  end

  return data_out
end

function multiply_dimensionwise_squeezed_avx!(
    data_out::AbstractArray{<:Any, 4}, data_in::AbstractArray{<:Any, 4}, vandermonde)
  n_vars      = size(data_out, 1)
  n_nodes_out = size(vandermonde, 1)
  n_nodes_in  = size(vandermonde, 2)
  data_out .= zero(eltype(data_out))

  @boundscheck begin
    inbounds = (size(data_out, 2) == size(data_out, 3) == size(data_out, 4) == n_nodes_out) &&
               (size(data_in,  1) == n_vars) &&
               (size(data_in,  2) == size(data_in,  3) == size(data_in,  4) == n_nodes_in)
    inbounds || throw(BoundsError())
  end

  @avx for k in 1:n_nodes_out, j in 1:n_nodes_out, i in 1:n_nodes_out
    for v in 1:n_vars
      acc = zero(eltype(data_out))
      for kk in 1:n_nodes_in, jj in 1:n_nodes_in, ii in 1:n_nodes_in
        acc += vandermonde[i, ii] * vandermonde[j, jj] * vandermonde[k, kk] * data_in[v, ii, jj, kk]
      end
      data_out[v, i, j, k] = acc
    end
  end

  return data_out
end

function multiply_dimensionwise_squeezed_tullio!(
    data_out::AbstractArray{<:Any, 4}, data_in::AbstractArray{<:Any, 4}, vandermonde)

  @tullio threads=false data_out[v, i, j, k] = vandermonde[i, ii] * vandermonde[j, jj] * vandermonde[k, kk] * data_in[v, ii, jj, kk]

  return data_out
end

function run_benchmarks_3d(n_vars=4, n_nodes_in=4, n_nodes_out=2*n_nodes_in)
  data_in  = randn(n_vars, n_nodes_in,  n_nodes_in,  n_nodes_in)
  data_out = randn(n_vars, n_nodes_out, n_nodes_out, n_nodes_out)
  vandermonde_dynamic = randn(n_nodes_out, n_nodes_in)
  vandermonde_static  = SMatrix{n_nodes_out, n_nodes_in}(vandermonde_dynamic)

  println("\n\n# 3D   ", "#"^70)
  println("n_vars      = ", n_vars)
  println("n_nodes_in  = ", n_nodes_in)
  println("n_nodes_out = ", n_nodes_out)
  println()

  println("\n","multiply_dimensionwise_sequential!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_sequential!(data_out, data_in, vandermonde_dynamic)
  data_out_copy = copy(data_out)
  display(@benchmark multiply_dimensionwise_sequential!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_sequential!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential!($(data_out), $(data_in), $(vandermonde_static)))
  println()

  println("\n","multiply_dimensionwise_sequential_avx!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_sequential_avx!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_avx!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_sequential_avx!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_avx!($(data_out), $(data_in), $(vandermonde_static)))
  println()

  println("\n","multiply_dimensionwise_sequential_tullio!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_sequential_tullio!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_sequential_tullio!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_static)))
  println()

  println("\n","multiply_dimensionwise_sequential_nexpr!")
  println("vandermonde_static")
  multiply_dimensionwise_sequential_nexpr!(data_out, data_in, vandermonde_static, Val(n_vars))
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_sequential_nexpr!($(data_out), $(data_in), $(vandermonde_static), $(Val(n_vars))))
  println()

  println("\n","multiply_dimensionwise_squeezed!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_squeezed!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_squeezed!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed!($(data_out), $(data_in), $(vandermonde_static)))
  println()

  println("\n","multiply_dimensionwise_squeezed_avx!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_squeezed_avx!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_avx!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_squeezed_avx!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_avx!($(data_out), $(data_in), $(vandermonde_static)))
  println()

  println("\n","multiply_dimensionwise_squeezed_tullio!")
  println("vandermonde_dynamic")
  multiply_dimensionwise_squeezed_tullio!(data_out, data_in, vandermonde_dynamic)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_tullio!($(data_out), $(data_in), $(vandermonde_dynamic)))
  println("\n", "vandermonde_static")
  multiply_dimensionwise_squeezed_tullio!(data_out, data_in, vandermonde_static)
  @assert data_out ≈ data_out_copy
  display(@benchmark multiply_dimensionwise_squeezed_tullio!($(data_out), $(data_in), $(vandermonde_static)))
  println()

  nothing
end

# TODO
run_benchmarks_3d(5, 4, 4) # n_vars, n_nodes_in, n_nodes_out

function compute_benchmarks_3d(n_vars, n_nodes_in, n_nodes_out)
  data_in  = randn(n_vars, n_nodes_in,  n_nodes_in,  n_nodes_in)
  data_out = randn(n_vars, n_nodes_out, n_nodes_out, n_nodes_out)
  vandermonde_dynamic = randn(n_nodes_out, n_nodes_in)
  vandermonde_static  = SMatrix{n_nodes_out, n_nodes_in}(vandermonde_dynamic)
  tmp1 = zeros(eltype(data_out), n_vars, n_nodes_out, n_nodes_in,  n_nodes_in)
  tmp2 = zeros(eltype(data_out), n_vars, n_nodes_out, n_nodes_out, n_nodes_in)

  println("n_vars = ", n_vars, "; n_nodes_in = ", n_nodes_in, "; n_nodes_out = ", n_nodes_out)
  sequential_dynamic = @benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_dynamic))
  sequential_static  = @benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_static))
  #FIXME sequential_nexpr   = @benchmark multiply_dimensionwise_sequential_nexpr!($(data_out), $(data_in), $(vandermonde_static), $(Val(n_vars)))
  sequential_dynamic_prealloc = @benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_dynamic), $(tmp1), $(tmp2))
  sequential_static_prealloc  = @benchmark multiply_dimensionwise_sequential_tullio!($(data_out), $(data_in), $(vandermonde_static),  $(tmp1), $(tmp2))
  squeezed_dynamic = @benchmark multiply_dimensionwise_squeezed_tullio!($(data_out), $(data_in), $(vandermonde_dynamic))
  squeezed_static  = @benchmark multiply_dimensionwise_squeezed_tullio!($(data_out), $(data_in), $(vandermonde_static))

  return time(median(sequential_dynamic)),
         time(median(sequential_static)),
         NaN, #FIXME time(median(sequential_nexpr)),
         time(median(sequential_dynamic_prealloc)),
         time(median(sequential_static_prealloc)),
         time(median(squeezed_dynamic)),
         time(median(squeezed_static))
end

function compute_benchmarks_3d(n_vars_list, n_nodes_in_list)
  sequential_dynamic          = zeros(length(n_vars_list), length(n_nodes_in_list))
  sequential_static           = zeros(length(n_vars_list), length(n_nodes_in_list))
  sequential_nexpr            = zeros(length(n_vars_list), length(n_nodes_in_list))
  sequential_dynamic_prealloc = zeros(length(n_vars_list), length(n_nodes_in_list))
  sequential_static_prealloc  = zeros(length(n_vars_list), length(n_nodes_in_list))
  squeezed_dynamic            = zeros(length(n_vars_list), length(n_nodes_in_list))
  squeezed_static             = zeros(length(n_vars_list), length(n_nodes_in_list))

  # n_nodes_out = n_nodes_in
  # mortar
  # superset of n_vars = 1, n_nodes_out = n_nodes_in, used for blending
  for (idx_nodes, n_nodes_in) in enumerate(n_nodes_in_list)
    for (idx_variables, n_vars) in enumerate(n_vars_list)
      n_nodes_out = n_nodes_in
      sequential_dynamic[idx_variables, idx_nodes],
      sequential_static[idx_variables, idx_nodes],
      sequential_nexpr[idx_variables, idx_nodes],
      sequential_dynamic_prealloc[idx_variables, idx_nodes],
      sequential_static_prealloc[idx_variables, idx_nodes],
      squeezed_dynamic[idx_variables, idx_nodes],
      squeezed_static[idx_variables, idx_nodes] =
        compute_benchmarks_3d(n_vars, n_nodes_in, n_nodes_out)
    end
  end
  BSON.@save "3D_nVarTotal_nNodesIn.bson" n_vars_list n_nodes_in_list sequential_dynamic sequential_static sequential_nexpr sequential_dynamic_prealloc sequential_static_prealloc squeezed_dynamic squeezed_static

  # TODO deactivated to save some time
  # # n_nodes_out = 2*n_nodes_in
  # # visualization
  # title = "n_vars = 2*n_vars, n_nodes_out = n_nodes_in"
  # for (idx_nodes, n_nodes_in) in enumerate(n_nodes_in_list)
  #   for (idx_variables, n_vars) in enumerate(n_vars_list)
  #     n_nodes_out = 2 * n_nodes_in
  #     sequential_dynamic[idx_variables, idx_nodes],
  #     sequential_static[idx_variables, idx_nodes],
  #     sequential_nexpr[idx_variables, idx_nodes],
  #     sequential_dynamic_prealloc[idx_variables, idx_nodes],
  #     sequential_static_prealloc[idx_variables, idx_nodes],
  #     squeezed_dynamic[idx_variables, idx_nodes],
  #     squeezed_static[idx_variables, idx_nodes] =
  #       compute_benchmarks_3d(n_vars, n_nodes_in, n_nodes_out)
  #   end
  # end
  # BSON.@save "3D_nVarTotal_2nNodesIn.bson" n_vars_list n_nodes_in_list sequential_dynamic sequential_static sequential_nexpr sequential_dynamic_prealloc sequential_static_prealloc squeezed_dynamic squeezed_static

  return nothing
end

# TODO
compute_benchmarks_3d(1:10, 2:10)


text:
# Autogenerated wrapper script for Xorg_libXmu_jll for x86_64-linux-musl
export libXmuu, libXmu

using Xorg_libXext_jll
using Xorg_libXt_jll
## Global variables
PATH = ""
LIBPATH = ""
LIBPATH_env = "LD_LIBRARY_PATH"

# Relative path to `libXmuu`
const libXmuu_splitpath = ["lib", "libXmuu.so"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libXmuu_path = ""

# libXmuu-specific global declaration
# This will be filled out by __init__()
libXmuu_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libXmuu = "libXmuu.so.1"

# Relative path to `libXmu`
const libXmu_splitpath = ["lib", "libXmu.so"]

# This will be filled out by __init__() for all products, as it must be done at runtime
libXmu_path = ""

# libXmu-specific global declaration
# This will be filled out by __init__()
libXmu_handle = C_NULL

# This must be `const` so that we can use it with `ccall()`
const libXmu = "libXmu.so.6"

"""
Open all libraries
"""
function __init__()
    global artifact_dir = abspath(artifact"Xorg_libXmu")

    # Initialize PATH and LIBPATH environment variable listings
    global PATH_list, LIBPATH_list
    # We first need to add to LIBPATH_list the libraries provided by Julia
    append!(LIBPATH_list, [joinpath(Sys.BINDIR, Base.LIBDIR, "julia"), joinpath(Sys.BINDIR, Base.LIBDIR)])
    # From the list of our dependencies, generate a tuple of all the PATH and LIBPATH lists,
    # then append them to our own.
    foreach(p -> append!(PATH_list, p), (Xorg_libXext_jll.PATH_list, Xorg_libXt_jll.PATH_list,))
    foreach(p -> append!(LIBPATH_list, p), (Xorg_libXext_jll.LIBPATH_list, Xorg_libXt_jll.LIBPATH_list,))

    global libXmuu_path = normpath(joinpath(artifact_dir, libXmuu_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libXmuu_handle = dlopen(libXmuu_path)
    push!(LIBPATH_list, dirname(libXmuu_path))

    global libXmu_path = normpath(joinpath(artifact_dir, libXmu_splitpath...))

    # Manually `dlopen()` this right now so that future invocations
    # of `ccall` with its `SONAME` will find this path immediately.
    global libXmu_handle = dlopen(libXmu_path)
    push!(LIBPATH_list, dirname(libXmu_path))

    # Filter out duplicate and empty entries in our PATH and LIBPATH entries
    filter!(!isempty, unique!(PATH_list))
    filter!(!isempty, unique!(LIBPATH_list))
    global PATH = join(PATH_list, ':')
    global LIBPATH = join(LIBPATH_list, ':')

    # Add each element of LIBPATH to our DL_LOAD_PATH (necessary on platforms
    # that don't honor our "already opened" trick)
    #for lp in LIBPATH_list
    #    push!(DL_LOAD_PATH, lp)
    #end
end  # __init__()

ext: jl


text:
using OptimizationDynamics
const iLQR = OptimizationDynamics.IterativeLQR
using LinearAlgebra
using Random

# ## visualization
vis = Visualizer()
render(vis)

# ## mode
MODE = :friction
MODE = :frictionless

# ## state-space model
h = 0.05
T = 51

if MODE == :friction
	im_dyn = ImplicitDynamics(cartpole_friction, h, eval(r_cartpole_friction_func), eval(rz_cartpole_friction_func), eval(rθ_cartpole_friction_func);
		r_tol=1.0e-8, κ_eval_tol=1.0e-4, κ_grad_tol=1.0e-3, no_impact=true)
        cartpole_friction.friction .= [0.35; 0.35]
        ## cartpole_friction.friction .= [0.25; 0.25]
        ## cartpole_friction.friction .= [0.1; 0.1]
        ## cartpole_friction.friction .= [0.01; 0.01]
else
	im_dyn = ImplicitDynamics(cartpole_frictionless, h, eval(r_cartpole_frictionless_func), eval(rz_cartpole_frictionless_func), eval(rθ_cartpole_frictionless_func);
    	        r_tol=1.0e-8, κ_eval_tol=1.0, κ_grad_tol=1.0, no_impact=true, no_friction=true)
end

nx = 2 * cartpole_friction.nq
nu = cartpole_friction.nu

# ## iLQR model
ilqr_dyn = iLQR.Dynamics((d, x, u, w) -> f(d, im_dyn, x, u, w),
					(dx, x, u, w) -> fx(dx, im_dyn, x, u, w),
					(du, x, u, w) -> fu(du, im_dyn, x, u, w),
					nx, nx, nu)
model = [ilqr_dyn for t = 1:T-1];

# ## initial conditions
q0 = [0.0; 0.0]
q1 = [0.0; 0.0]
qT = [0.0; π]
q_ref = [0.0; π]

x1 = [q1; q1]
xT = [qT; qT]

# ## objective
function objt(x, u, w)
	J = 0.0
	J += transpose(u) * u
	return J
end

function objT(x, u, w)
        J = 0.0
        J += transpose(x - xT) * (x - xT)
        return J
end

ct = iLQR.Cost(objt, nx, nu)
cT = iLQR.Cost(objT, nx, 0)
obj = [[ct for t = 1:T-1]..., cT];

# ## constraints
function terminal_con(x, u, w)
    [
    x - xT; # goal
    ]
end

cont = iLQR.Constraint()
conT = iLQR.Constraint(terminal_con, nx, 0)
cons = [[cont for t = 1:T-1]..., conT];

# ## rollout
ū = [(t == 1 ? -1.5 : 0.0) * ones(nu) for t = 1:T-1] # set value to -1.0 when friction coefficient = 0.25
x̄ = iLQR.rollout(model, x1, ū)
q̄ = state_to_configuration(x̄)
visualize!(vis, cartpole_friction, q̄, Δt=h);

# ## solver
solver = iLQR.solver(model, obj, cons,
    opts=iLQR.Options(linesearch = :armijo,
    α_min=1.0e-5,
    obj_tol=1.0e-5,
    grad_tol=1.0e-3,
    max_iter=100,
    max_al_iter=20,
    con_tol=0.005,
    ρ_init=1.0,
    ρ_scale=10.0,
    verbose=false))
iLQR.initialize_controls!(solver, ū)
iLQR.initialize_states!(solver, x̄);

# ## solve
iLQR.reset!(solver.s_data)
@time iLQR.solve!(solver);

@show iLQR.eval_obj(solver.m_data.obj.costs, solver.m_data.x, solver.m_data.u, solver.m_data.w)
@show solver.s_data.iter[1]
@show norm(terminal_con(solver.m_data.x[T], zeros(0), zeros(0)), Inf)
@show solver.s_data.obj[1] # augmented Lagrangian cost

# ## solution
x_sol, u_sol = iLQR.get_trajectory(solver)
q_sol = state_to_configuration(x_sol)
visualize!(vis, cartpole_friction, q_sol, Δt=h);

# ## benchmark
solver.options.verbose = false
@benchmark iLQR.solve!($solver, x̄, ū) setup=(x̄=deepcopy(x̄), ū=deepcopy(ū));


text:
"""
ksintegrate: integrate kuramoto-sivashinsky equation (Julia)
       u_t = -u*u_x - u_xx - u_xxxx, domain x in [0,Lx], periodic BCs

 inputs
          u = initial condition (vector of u(x) values on uniform gridpoints))
         Lx = domain length
         dt = time step
         Nt = number of integration timesteps
      nsave = save every nsave-th time step

 outputs

          u = final state, vector of u(x, Nt*dt) at uniform x gridpoints

This implementation has two improvements over ksintegrateNaive.jl. It uses
   (1) in-place FFTs
   (2) loop fusion: Julia can translate arithmetic vector expressions in dot syntax
       to single for loop over the components, which should be much faster than
       constructing a temporary vector for each operation in the vector expression.
"""
function ksintegrateInplace(u, Lx, dt, Nt)
    u = (1+0im)*u                       # force u to be complex
    Nx = length(u)                      # number of gridpoints
    kx = vcat(0:Nx/2-1, 0:0, -Nx/2+1:-1)# integer wave #s,  exp(2*pi*i*kx*x/L)
    alpha = 2*pi*kx/Lx                  # real wavenumbers, exp(i*alpha*x)
    D = 1im*alpha                       # spectral D = d/dx operator
    L = alpha.^2 - alpha.^4             # spectral L = -D^2 - D^4 operator
    G = -0.5*D                          # spectral -1/2 D operator, to eval -u u_x = 1/2 d/dx u^2

    # convenience variables
    dt2  = dt/2
    dt32 = 3*dt/2
    A_inv = (ones(Nx) - dt2*L).^(-1)
    B     =  ones(Nx) + dt2*L

    # compute FFTW plans
    FFT! = plan_fft!(u, flags=FFTW.ESTIMATE)
    IFFT! = plan_ifft!(u, flags=FFTW.ESTIMATE)

    # compute nonlinear term Nn == -u u_x and Nuprev (Nu at prev timestep)
    Nn  = G.*fft(u.^2)  # Nn == -1/2 d/dx (u^2) = -u u_x
    Nn1 = copy(Nn)      # use Nn1 = Nn at first time step
    FFT!*u

    # timestepping loop
    for n = 0:Nt

        Nn1 .= Nn   # shift nonlinear term in time
        Nn .= u        # put u into N in prep for comp of nonlineat

        IFFT!*Nn       # transform Nn to gridpt values, in place
        Nn .= Nn.*Nn   # collocation calculation of u^2
        FFT!*Nn        # transform Nn back to spectral coeffs, in place

        Nn .= G.*Nn

        # loop fusion! Julia translates the folling line of code to a single for loop.
        u .= A_inv .*(B .* u .+ dt32.*Nn .- dt2.*Nn1)
    end

    IFFT!*u
    real(u)
end

max_stars_repo_name: raphbacher/julia-intro

max_issues_repo_name: raphbacher/julia-intro

max_forks_repo_name: raphbacher/julia-intro


text:
mutable struct RatePseudoRootJacobianAllElements
  pseudoRoot::Matrix{Float64}
  aliveIndex::Int
  taus::Vector{Float64}
  pseudoBumps::Vector{Matrix{Float64}}
  displacements::Vector{Float64}
  factors::Int
  e::Matrix{Float64}
  ratios::Vector{Float64}
end

function RatePseudoRootJacobianAllElements(pseudoRoot::Matrix{Float64},
                                            aliveIndex::Int,
                                            numeraire::Int,
                                            taus::Vector{Float64},
                                            displacements::Vector{Float64})
  pr_rows, factors = size(pseudoRoot)
  e = Matrix{Float64}(undef, pr_rows, factors)
  numberRates = length(taus)
  ratios = Vector{Float64}(undef, numberRates)

  aliveIndex == numeraire || error("we can only do discrete compounding MM account so alive Index must equal numeraire")
  pr_rows == numberRates || error("pseudoRoot rows != taus size")
  length(displacements) == numberRates || error("displacements size != taus size")

  return RatePseudoRootJacobianAllElements(pseudoRoot, aliveIndex, taus, Vector{Matrix{Float64}}(), displacements, factors, e, ratios)
end

function get_bumps!(ratePseudo::RatePseudoRootJacobianAllElements, oldRates::Vector{Float64}, discountRatios::Vector{Float64}, newRates::Vector{Float64},
                    gaussians::Vector{Float64}, B::Vector{Matrix{Float64}})
  numberRates = length(ratePseudo.taus)

  length(B) == numberRates || error("B must be the same size as numberRates")
  @simd for j in eachindex(B)
    @inbounds size(B[j]) == (numberRates, ratePseudo.factors) || error("matrix in B not of proper size")
  end

  @simd for j = ratePseudo.aliveIndex:numberRates
    @inbounds ratePseudo.ratios[j] = (oldRates[j] + ratePseudo.displacements[j]) * discountRatios[j+1]
  end

  @inbounds @simd for f = 1:ratePseudo.factors
    ratePseudo.e[ratePseudo.aliveIndex, f] = 0.0
    for j = ratePseudo.aliveIndex+1:numberRates
      ratePseudo.e[j, f] = ratePseudo.e[j-1, f] + ratePseudo.ratios[j-1] * ratePseudo.pseudoRoot[j-1, f]
    end
  end

  # nullify B for rates that have already reset
  for j = 1:ratePseudo.aliveIndex, k = 1:numberRates, f = 1:ratePseudo.factors
    @inbounds B[j][k, f] = 0.0
  end

  @inbounds for f = 1:ratePseudo.factors, j = ratePseudo.aliveIndex:numberRates
    @simd for k = ratePseudo.aliveIndex:j
      B[j][k, f] = newRates[j] * ratePseudo.ratios[k] * ratePseudo.taus[k] * ratePseudo.pseudoRoot[j, f]
    end

    tmp = 2.0 * ratePseudo.ratios[j] * ratePseudo.taus[j] * ratePseudo.pseudoRoot[j, f]
    tmp -= ratePseudo.pseudoRoot[j, f]
    tmp += ratePseudo.e[j, f] * ratePseudo.taus[j]
    tmp += gaussians[f]
    tmp *= (newRates[j] + ratePseudo.displacements[j])

    B[j][j, f] = tmp

    for k = 1:ratePseudo.aliveIndex
      B[j][k, f] = 0.0
    end

    for k = j+1:numberRates
      B[j][k, f] = 0.0
    end
  end

  return ratePseudo, B
end


text:
using Test

using MarkovChains
@testset "inftime_solve" begin
@testset "acyclic1" begin
    chain = ContMarkovChain()
    n1 = add_state!(chain)
    n2 = add_state!(chain)
    n3 = add_state!(chain)
    add_transition!(chain, n1, n2, 2.0)
    add_transition!(chain, n2, n3, 4.0)
    init_prob = fill(0.0, state_count(chain))
    init_prob[1] = 1.0
    res = inftime_solve(chain, init_prob)
    cumtime = collect(map(state -> state_cumtime(res, state), n1:n3))
    @test cumtime[n1] ≈ 0.5
    @test cumtime[n2] ≈ 0.25
    @test cumtime[n3] ≈ Inf
    prob = collect(map(state -> state_prob(res, state), n1:n3))
    @test prob[n1] ≈ 0
    @test prob[n2] ≈ 0
    @test prob[n3] ≈ 1

    @test mean_time_to_absorption(chain, init_prob) == 0.75
end
@testset "irreducible1" begin
    chain = ContMarkovChain()
    n0 = add_state!(chain)
    n1 = add_state!(chain)
    n2 = add_state!(chain)
    n3 = add_state!(chain)
    add_transition!(chain, n0, n1, 1.0)
    add_transition!(chain, n1, n2, 1.0)
    add_transition!(chain, n2, n3, 1.0)
    add_transition!(chain, n3, n2, 3.0)
    add_transition!(chain, n2, n1, 2.0)
    add_transition!(chain, n1, n0, 1.0)
    Q = trans_rate_matrix(chain)
    init_prob = fill(0.0, state_count(chain))
    init_prob[1] = 0.1
    init_prob[2] = 0.9
    res = inftime_solve(chain, init_prob)
    cumtime = collect(map(state -> state_cumtime(res, state), n0:n3))
    for t in cumtime
        @test t == Inf
    end
    prob = collect(map(state -> state_prob(res, state), n0:n3))
    @test prob ≈ [0.375, 0.375, 0.1875, 0.0625]
end
@testset "multiple_comps" begin
    chain = ContMarkovChain()
    t0 = add_state!(chain)
    t1 = add_state!(chain)
    t2 = add_state!(chain)
    m1 = add_state!(chain)
    m2 = add_state!(chain)
    n1 = add_state!(chain)
    n2 = add_state!(chain)
    add_transition!(chain, t0, t1, 1.0)
    add_transition!(chain, t1, t2, 1.0)
    add_transition!(chain, t2, t1, 2.0)
    add_transition!(chain, t1, m1, 1.0)
    add_transition!(chain, t2, n1, 3.0)
    add_transition!(chain, m1, m2, 2.0)
    add_transition!(chain, m2, m1, 1.0)
    add_transition!(chain, n1, n2, 2.0)
    add_transition!(chain, n2, n1, 1.0)
    init_prob = fill(0.0, state_count(chain))
    init_prob[1] = 1.0
    res = inftime_solve(chain, init_prob)
    cumtime = collect(map(state -> state_cumtime(res, state), t0:n2))
    for t in cumtime[[m1, m2, n1, n2]]
        @test t == Inf
    end
    for t in cumtime[[t0, t1, t2]]
        @test t > 0
    end
    prob = collect(map(state -> state_prob(res, state), t0:n2))
    for p in prob[[t0, t1, t2]]
        @test p == 0.0
    end
    for p in prob[[m1, m2, n1, n2]]
        @test p > 0
    end

    @test isinf(mean_time_to_absorption(chain, init_prob))
end
using SparseArrays
@testset "multiple_comps_init_bottom" begin
    chain = ContMarkovChain()
    t0 = add_state!(chain)
    t1 = add_state!(chain)
    t2 = add_state!(chain)
    m1 = add_state!(chain)
    m2 = add_state!(chain)
    n1 = add_state!(chain)
    n2 = add_state!(chain)
    add_transition!(chain, t0, t1, 1.0)
    add_transition!(chain, t1, t2, 1.0)
    add_transition!(chain, t2, t1, 2.0)
    add_transition!(chain, t1, m1, 1.0)
    add_transition!(chain, t2, n1, 3.0)
    add_transition!(chain, m1, m2, 2.0)
    add_transition!(chain, m2, m1, 1.0)
    add_transition!(chain, n1, n2, 2.0)
    add_transition!(chain, n2, n1, 1.0)
    init_prob = sparsevec([n2, n1, m1], [0.7, 0.2, 0.1])
    res = inftime_solve(chain, init_prob)
    cumtime = collect(map(state -> state_cumtime(res, state), t0:n2))
    for t in cumtime[[m1, m2, n1, n2]]
        @test t == Inf
    end
    for t in cumtime[[t0, t1, t2]]
        @test t == 0.0
    end
    prob = collect(map(state -> state_prob(res, state), t0:n2))
    for p in prob[[t0, t1, t2]]
        @test p == 0.0
    end
    @test prob[n1] == 1.0 / 3.0 * 0.9
    @test prob[n2] == 2.0 / 3.0 * 0.9
    @test prob[m1] == 1.0 / 3.0 * 0.1
    @test prob[m2] == 2.0 / 3.0 * 0.1
end
@testset "one_state" begin
    chain = ContMarkovChain()
    t0 = add_state!(chain)
    init_prob = sparsevec([t0], [1.0])
    res = inftime_solve(chain, init_prob)
    @test state_cumtime(res, t0) == Inf
    @test state_prob(res, t0) == 1.0
    @test mean_time_to_absorption(chain, init_prob) == 0.0
end
@testset "zero_state" begin
    chain = ContMarkovChain()
    init_prob = sparsevec([], [])
    res = inftime_solve(chain, init_prob)
end

end


text:
export sbesselj, shankelh1, diffsbessel, diffbessel
export diffbesselj, diffhankelh1, diffsbesselj, diffshankelh1
export gaunt_coefficient
export associated_legendre_indices, spherical_harmonics_indices, lm_to_spherical_harmonic_index
export spherical_harmonics
export cartesian_to_radial_coordinates, radial_to_cartesian_coordinates
export atan

#NOTE spherical bessel and hankel functions soon coming to SpecialFunctions.jl https://github.com/JuliaMath/SpecialFunctions.jl/pull/196

#NOTE atan(x,y) is defined here but will likely be added to base soon.

import Base.atan
atan(y::Complex,x::Complex) = - im * log( (x+y*im) / sqrt(x^2+y^2) )

"""
    sbesselj(m,x)

Returns the spherical besselj function. The order is 'm' and the argument is 'x'. Note 'x' can be a complex number.
"""
# function sbesselj(m::Number,x::T) where T<: Union{F,Complex{F}} where F <: AbstractFloat
    # if (abs(x) > eps(F))
    #     return sqrt(pi/(T(2)*x)) * besselj(m+1/2,x)
    # else
    #     return (m > 0 ? zero(T) : one(T))
    # end
# end
function sbesselj(nu::Number, x::T) where {T}
    besselj_nuhalf_x = besselj(nu + one(nu)/2, x)
    if abs(x) ≤ sqrt(eps(real(zero(besselj_nuhalf_x))))
        nu == 0 ? one(besselj_nuhalf_x) : zero(besselj_nuhalf_x)
    else
        √((float(T))(π)/2x) * besselj_nuhalf_x
    end
end

"""
    shankelh1(m,x)

Returns the spherical hankel function of the first kind. The order is 'm' and the argument is 'x'. Note 'x' can be a complex number.
"""
shankelh1(nu::Number,x::T) where T = √((float(T))(π)/2x) * hankelh1(nu+one(nu)/2,x)

"""
    diffsbessel(f::Function,m,x)

Differentiates the spherical bessel function 'f' (for any spherical bessel). The order is 'm' and the argument is 'x'. Note 'x' can be a complex number.
"""
function diffsbessel(f::Function,n::Number,z::Number)
    return f(n-1,z) - (n+1) * f(n,z) / z
end

function diffsbesselj(n::Number,z::Number)
    return if n == 0
        - sbesselj(1,z) # case due to numerical stability
    else
        sbesselj(n-1,z) - (n+1) * sbesselj(n,z) / z
    end
end

function diffshankelh1(n::Number,z::Number)
    return shankelh1(n-1,z) - (n+1) * shankelh1(n,z) / z
end

"""
    diffbessel(f::Function,m,x,n::Int)

Differentiates 'n' times any bessel function 'f' of order 'm' and at the argument 'x'.
"""
function diffbessel(f::Function,m::Number,z,n::Int)
    if n == 0
        return f(m, z)
    elseif n > 0
        n = n - 1
        return 0.5*(diffbessel(f,m-1,z,n) - diffbessel(f,m+1,z,n))
    else
        error("Can not differentiate a negative number of times")
    end
end

"""
    diffhankelh1(m,x,n::Int)

Differentiates 'n' times the hankelh1 function of order 'm' and at the argument 'x'.
"""
function diffhankelh1(m::Number,z::T,n::Int) where T<:Number
    if n == 0
        return hankelh1(m, z)
    elseif n > 0
        n = n - 1
        return 0.5*(diffhankelh1(m-1,z,n) - diffhankelh1(m+1,z,n))
    else
        error("Can not differentiate a negative number of times")
    end
end

"Derivative of Hankel function of the first kind"
diffhankelh1(n,z) = 0.5*(hankelh1(-1 + n, z) - hankelh1(1 + n, z))

"Derivative of Bessel function of first kind"
diffbesselj(n::Number,z::T) where T<:Number = 0.5*(besselj(-1 + n, z) - besselj(1 + n, z))

"m-th Derivative of Hankel function of the first kind"
function diffbesselj(n::Number,z::T,m::Int) where T<:Number
    if m == 0
        return besselj(n, z)
    elseif m > 0
        m = m - 1
        return T(0.5)*(diffbesselj(n-1,z,m) - diffbesselj(n+1,z,m))
    else
        error("Can not differentiate a negative number of times")
    end
end

"""
    gaunt_coefficient(l1,m1,l2,m2,l3,m3)

A version of the Gaunt coefficients which are used to write the product of two spherical harmonics. If Y_{l,m} is a complex spherical harmonic, with the typical phase conventions from quantum mechanics, then:

    gaunt_coefficient(l1,m1,l2,m2,l3,m3) = 4*π*im^{l2+l3-l1} Integral[Y_{l1,m1}*conj(Y_{l2,m2})*conj(Y_{l3,m3})]

where the integral is over the solid angle.

The most standard gaunt coefficients `G(l1,m1;l2,m2;l3)` are related through the identity:

    4pi * G(l1,m1;l2,m2;l3) = im^(l1-l2-l3) * (-1)^m2 * gaunt_coefficient(l1,m1,l2,-m2,l3,m1+m2)

"""
function gaunt_coefficient(T::Type{<:AbstractFloat},l1::Int,m1::Int,l2::Int,m2::Int,l3::Int,m3::Int)
    # note the wigner3j has only one convention, and is highly symmetric.
    return (one(T)*im)^(l2+l3-l1) * (-T(1))^m1 * sqrt(4pi*(2*l1+1)*(2*l2+1)*(2*l3+1)) *
        wigner3j(T,l1,l2,l3,0,0,0) * wigner3j(T,l1,l2,l3,m1,-m2,-m3)
end
gaunt_coefficient(l1::Int,m1::Int,l2::Int,m2::Int,l3::Int,m3::Int) = gaunt_coefficient(Float64,l1,m1,l2,m2,l3,m3)

lm_to_spherical_harmonic_index(l::Int,m::Int)::Int = l^2 + m + l + 1

function spherical_harmonics_indices(l_max::Int)
    ls = [l for l in 0:l_max for m in -l:l]
    ms = [m for l in 0:l_max for m in -l:l]

    return ls, ms
end

function associated_legendre_indices(l_max::Int)
    ls = [l for l in 0:l_max for m in 0:l]
    ms = [m for l in 0:l_max for m in 0:l]

    return ls, ms
end

"""
`spherical_harmonics(l_max::Int, θ::T, φ::T)`

returns a vector of all spherical harmonics with degree `l <= l_max`. The degree and order (indices) of the elements of the vector are given by `spherical_harmonics_indices(l_max::Int)`.

The associated legendre polynomials are taken from the package GSL.jl.
"""
function spherical_harmonics(l_max::Int, θ::Complex{T}, φ::Union{T,Complex{T}}) where T <: AbstractFloat

    throw(DomainError(θ, "Currently GLS.jl is used to calculate associated legendre polynomials, which only take real angles as arguments. Hopefully soon SpecialFunctions.jl will implement associated Legendre for complex arguments: https://github.com/JuliaMath/SpecialFunctions.jl/pull/175"))

end
function spherical_harmonics(l_max::Int, θ::T, φ::T) where T <: AbstractFloat

    ls, ms = associated_legendre_indices(l_max)
    Plm_arr = sf_legendre_array(GSL_SF_LEGENDRE_SPHARM, l_max, cos(θ))[eachindex(ls)]

    Ylm_vec = Vector{Complex{T}}(undef, (l_max+1)^2)
    Ylm_vec[1] = Plm_arr[1]

    ind1 = 1
    ind2 = 1
    for i = 1:l_max
        inds1 = (ind1+i):(ind1+2i)
        Ylm_vec[(ind2+i):(ind2+3i)] = [reverse((-1).^ms[inds1[2:end]] .* conj(Plm_arr[inds1[2:end]])); Plm_arr[inds1]]
        ind1 += i
        ind2 += 2i
    end

    ls, ms = spherical_harmonics_indices(l_max)
    Ylm_vec = (-one(T)) .^ ms .* exp.(ms .* (im*φ)) .* Ylm_vec

    return Ylm_vec
end

cartesian_to_radial_coordinates(x::Vector) = cartesian_to_radial_coordinates(SVector(x...))
radial_to_cartesian_coordinates(θ::Vector) = radial_to_cartesian_coordinates(SVector(θ...))

function cartesian_to_radial_coordinates(x::SVector{3,CT}) where CT
    r = sqrt(sum(x .^2)) # note this is, and should be, a complex number when x is a complex vector
    θ = atan(sqrt(x[1]^2+x[2]^2),x[3])
    φ = (!iszero(x[1]) || !iszero(x[2])) ? atan(x[2], x[1]) : zero(CT)
    return [r,θ,φ]
end

function cartesian_to_radial_coordinates(x::SVector{2,CT}) where CT
    r = sqrt(sum(x .^2)) # note this should be complex if x is complex
    θ = atan(x[2], x[1])
    return [r,θ]
end

function radial_to_cartesian_coordinates(rθφ::SVector{3,CT}) where CT
    r, θ, φ = rθφ
    x = r * sin(θ) * cos(φ)
    y = r * sin(θ) * sin(φ)
    z = r * cos(θ)

    return [x,y,z]
end

function radial_to_cartesian_coordinates(rθ::SVector{2,CT}) where CT
    r, θ = rθ
    x = r * cos(θ)
    y = r * sin(θ)

    return [x,y]
end


text:
using TermWin

TermWin.initsession()
arr = map( x->utf8(x), readdir() )
v = newTwMultiSelect( rootTwScreen, arr, posy=:center, posx=:center, orderable=true, substrsearch=true )
v.title = "Input: "
activateTwObj( rootTwScreen )
ret = v.value
#unregisterTwObj( rootTwScreen, v )
TermWin.endsession()
println( "You chose ", string( ret ) )


text:
module Quon

export yang_baxter_param, yang_baxter_param_inv, change_direction, change_direction!
export PlanarMultigraph, Tait,
    vertices, faces, half_edges,
    src, dst, half_edge, face,
    next, prev, twin, α, ϕ, σ, σ_inv,
    nv, ne, nf, nhe, is_boundary,
    out_half_edge, surrounding_half_edge,
    trace_face, trace_vertex, neighbors,
    rem_vertex!, rem_edge!, rem_face!, merge_graph!,
    add_edge!, add_vertex!,
    check_faces, check_vertices, check_combinatorial_maps
export contract!, contract_boundary_vertices!,
    tensor_product!, merge_boundary_vertices!,
    quon_param, quon_params, genuses, is_genus, is_open_vertex
export planar_rx, planar_rz, planar_id, planar_copy, planar_zero_state, planar_swap,
    tait_rx, tait_rz, tait_id, tait_hadamard, tait_copy, tait_cz, tait_cnot,
    tait_zero_state, tait_swap

include("utils/yang_baxter.jl")

include("params/quon_const.jl")
include("params/quon_param.jl")

include("graphs/planar_multigraph.jl")
include("graphs/tait.jl")
include("graphs/predefined.jl")

export Rule, rewrite!, match!

include("match.jl")
include("rewrite.jl")
include("check.jl")
include("simplify.jl")

include("utils/zx.jl")
include("circuits.jl")

export push_gate!

export plot

include("plots.jl")

end


text:
export gdownload

const SPREADSHEET_PATTERN = "docs.google.com/spreadsheets"
const DRIVE_PATTERN = "drive.google.com/file/d"
const DOCS_PATTERN = "docs.google.com"

is_gsheet(url) = occursin(SPREADSHEET_PATTERN, url)
is_gfile(url) = occursin(DRIVE_PATTERN, url)
is_gdoc(url) = occursin(DOCS_PATTERN, url)

"""
    unshortlink(url)

return unshorten url or the url if it is not a short link
"""
function unshortlink(url; kw...)
    rq = HTTP.request("HEAD", url; redirect=false, status_exception=false, kw...)
    while !(is_gdoc(url) || is_gfile(url)) && (rq.status ÷ 100 == 3)
        url = HTTP.header(rq, "Location")
        rq = HTTP.request("HEAD", url; redirect=false, status_exception=false, kw...)
    end

    return url
end

function gsheet_handler(url; format=:csv)
    link, expo = splitdir(url)
    if startswith(expo, "edit") || (expo == "")
        url = link * "/export?format=$format"
    elseif startswith(expo, "export")
        url = replace(url, r"format=([a-zA-Z]*)(.*)"=>SubstitutionString("format=$format\\2"))
    end

    return url
end

function gfile_handler(url)
    # pattern of file path in google drive:
    # https://drive.google.com/file/d/<hash>/view?usp=sharing
    h = match(r"/file/d/([^\/]+)/", url)
    isnothing(h) && throw(ErrorException("Can't find google drive file ID in the url"))

    return "https://docs.google.com/uc?export=download&id=$(h.captures[])"
end

function find_gcode(cookies)
    for cookie in cookies
        (match(r"download_warning_", cookie.name) !== nothing) && (return cookie.value)
    end

    return
end

function find_filename(header)
    m = match(r"filename=\\\"(.*)\\\"", header)
    if isnothing(m)
        filename = "gdrive_downloaded-$(randstring())"
        @warn "File name not found, use `$filename`"
    else
        filename = m.captures[]
    end

    return filename
end

function download_gdrive(url, localdir)
    cookiejars = Dict{String, Set{HTTP.Cookies.Cookie}}()
    HTTP.request("GET", url; cookies=true, cookiejar=cookiejars)
    gcode = find_gcode(cookiejars["docs.google.com"])

    !isnothing(gcode) && (url = "$url&confirm=$gcode")

    filepath = Ref{String}("")
    HTTP.open(
        "GET", url, ["Range"=>"bytes=0-"],
        cookies=true, cookiejar=cookiejars, redirect_limit=10
    ) do stream
        response = HTTP.startread(stream)
        eof(stream) && return

        header = HTTP.header(response, "Content-Disposition")
        isempty(header) && return

        filepath[] = joinpath(localdir, find_filename(header))

        total_bytes = tryparse(Int64, rsplit(HTTP.header(response, "Content-Range"), '/'; limit=2)[end])
        isnothing(total_bytes) && (total_bytes = NaN)
        println("Total: $total_bytes bytes")

        downloaded_bytes = progress = 0
        print("Downloaded:\e[s")
        Base.open(filepath[], "w") do fh
            while !eof(stream)
                downloaded_bytes += write(fh, readavailable(stream))
                new_progress = 100downloaded_bytes ÷ total_bytes
                (new_progress > progress) && print("\e[u $downloaded_bytes bytes ($new_progress%)")
                progress = new_progress
            end
            println()
        end
    end

    return filepath[]
end

"""
    gdownload(url, localdir)

Download file or Google Sheet from Google drive.
"""
function gdownload(url, localdir)
    url = unshortlink(url)

    if is_gfile(url)
        url = gfile_handler(url)
    elseif is_gsheet(url)
        url = gsheet_handler(url)
    end

    is_gdoc(url) || throw("invalid url")
    download_gdrive(url, localdir)
end


text:
mutable struct SPMomentConstraintOutG <: SumProductRule{MomentConstraint} end
outboundType(::Type{SPMomentConstraintOutG}) = Message{GaussianWeightedMeanPrecision}
isApplicable(::Type{SPMomentConstraintOutG}, input_types::Vector{<:Type}) = (length(input_types) == 1) && (input_types[1] == Nothing)


text:
import RestrictedBoltzmannMachines as RBMs
import Zygote
using Test: @test, @testset
using Statistics: mean
using Random: bitrand
using RestrictedBoltzmannMachines: BinaryRBM, free_energy, ∂free_energy
using RestrictedBoltzmannMachines: Binary, Gaussian, ReLU, dReLU, pReLU, xReLU
using RestrictedBoltzmannMachines: ∂regularize!, ∂regularize, ∂regularize_fields

@testset "∂regularize!" begin
    rbm = BinaryRBM(randn(3,5), randn(3,2), randn(3,5,3,2))
    v = bitrand(3,5,100)
    vdims = ntuple(identity, ndims(rbm.visible))
    N = length(rbm.visible)

    l2_fields = rand()
    l1_weights = rand()
    l2_weights = rand()
    l2l1_weights = rand()

    gs = Zygote.gradient(rbm) do rbm
        F = mean(free_energy(rbm, v))
        L2_fields = sum(abs2, rbm.visible.θ)
        L1_weights = sum(abs, rbm.w)
        L2_weights = sum(abs2, rbm.w)
        L2L1_weights = sum(abs2, sum(abs, rbm.w; dims=vdims))
        return (
            F + l2_fields/2 * L2_fields + l1_weights * L1_weights +
            l2_weights/2 * L2_weights +
            l2l1_weights/(2N) * L2L1_weights
        )
    end

    ∂ = ∂free_energy(rbm, v)
    ∂regularize!(∂, rbm; l2_fields, l1_weights, l2_weights, l2l1_weights)

    @test only(gs).visible.θ ≈ ∂.visible.θ
    @test only(gs).hidden.θ ≈ ∂.hidden.θ
    @test only(gs).w ≈ ∂.w
end

@testset "∂regularize" begin
    rbm = BinaryRBM(randn(3,5), randn(3,2), randn(3,5,3,2))
    dims = ntuple(identity, ndims(rbm.visible))
    N = length(rbm.visible)

    l2_fields = rand()
    l1_weights = rand()
    l2_weights = rand()
    l2l1_weights = rand()

    gs = Zygote.gradient(rbm) do rbm
        return (
            l2_fields/2 * sum(abs2, rbm.visible.θ) +
            l1_weights * sum(abs, rbm.w) +
            l2_weights/2 * sum(abs2, rbm.w) +
            l2l1_weights/(2N) * sum(abs2, sum(abs, rbm.w; dims))
        )
    end

    ∂ = ∂regularize(rbm; l2_fields, l1_weights, l2_weights, l2l1_weights)

    @test only(gs).visible.θ ≈ ∂.visible.θ
    @test isnothing(only(gs).hidden) && isnothing(∂.hidden)
    @test only(gs).w ≈ ∂.w
end

@testset "∂regularize_fields" begin
    l2_fields = rand()

    layer = Binary(randn(3,5))
    gs = Zygote.gradient(layer) do layer
        l2_fields/2 * sum(abs2, layer.θ)
    end
    ∂ = ∂regularize_fields(layer; l2_fields)
    @test only(gs).θ ≈ ∂.θ

    layer = Gaussian(randn(3,5), rand(3,5))
    gs = Zygote.gradient(layer) do layer
        l2_fields/2 * sum(abs2, layer.θ)
    end
    ∂ = ∂regularize_fields(layer; l2_fields)
    @test only(gs).θ ≈ ∂.θ
    @test isnothing(only(gs).γ) && iszero(∂.γ)

    layer = dReLU(randn(3,5), randn(3,5), rand(3,5), rand(3,5))
    gs = Zygote.gradient(layer) do layer
        l2_fields/2 * (sum(abs2, layer.θp) + sum(abs2, layer.θn))
    end
    ∂ = ∂regularize_fields(layer; l2_fields)
    @test only(gs).θp ≈ ∂.θp
    @test only(gs).θn ≈ ∂.θn
    @test isnothing(only(gs).γp) && iszero(∂.γp)
    @test isnothing(only(gs).γn) && iszero(∂.γn)

    layer = pReLU(randn(3,5), rand(3,5), randn(3,5), rand(3,5) .- 0.5)
    gs = Zygote.gradient(layer) do layer
        l2_fields/2 * sum(abs2, layer.θ)
    end

    ∂ = ∂regularize_fields(layer; l2_fields)
    @test only(gs).θ ≈ ∂.θ
    @test isnothing(only(gs).γ) && iszero(∂.γ)
    @test isnothing(only(gs).Δ) && iszero(∂.Δ)
    @test isnothing(only(gs).η) && iszero(∂.η)

    layer = xReLU(randn(3,5), rand(3,5), randn(3,5), randn(3,5))
    gs = Zygote.gradient(layer) do layer
        l2_fields/2 * sum(abs2, layer.θ)
    end

    ∂ = ∂regularize_fields(layer; l2_fields)
    @test only(gs).θ ≈ ∂.θ
    @test isnothing(only(gs).γ) && iszero(∂.γ)
    @test isnothing(only(gs).Δ) && iszero(∂.Δ)
    @test isnothing(only(gs).ξ) && iszero(∂.ξ)
end


text:
export GoToDoor

using Random

mutable struct GoToDoor{W<:GridWorldBase} <: AbstractGridWorld
    world::W
    agent_pos::CartesianIndex{2}
    agent::Agent
end

function GoToDoor(;n=8, agent_start_pos=CartesianIndex(2,2), rng=Random.GLOBAL_RNG)
    objects = (EMPTY, WALL, (Door(c) for c in COLORS)...)
    world = GridWorldBase(objects, n, n)
    world[EMPTY, :, :] .= true
    world[WALL, [1,n], 1:n] .= true
    world[EMPTY, [1,n], 1:n] .= false
    world[WALL, 1:n, [1,n]] .= true
    world[EMPTY, 1:n, [1,n]] .= false

    door_pos = [(rand(rng, 2:n-1),1), (rand(rng, 2:n-1),n), (1,rand(rng, 2:n-1)), (n,rand(rng, 2:n-1))]
    door_colors = COLORS[randperm(rng, length(COLORS))][1:length(door_pos)]
    for (c, p) in zip(door_colors, door_pos)
        world[Door(c), p...] = true
        world[WALL, p...] = false
    end
    GoToDoor(world, agent_start_pos, Agent(dir=RIGHT))
end

function (w::GoToDoor)(::MoveForward)
    dir = get_dir(w.agent)
    dest = dir(w.agent_pos)
    if dest ∈ CartesianIndices((size(w.world, 2), size(w.world, 3))) && !w.world[WALL,dest]
        w.agent_pos = dest
    end
    w
end