text:
using Distributions
using Cytof5, Random, RCall
using JLD2, FileIO
import Printf
include("../sim_study/util.jl")
include("compress_data.jl")
Random.seed!(0)

function loadSingleObj(objPath)
  data = load(objPath)
  return data[collect(keys(data))[1]]
end

# TODO:
# Write a macro to reconstruct types

# Turn reconstructed type into type State
function convertState(s)
  println("Trying to convert reconstructed state to Cytof5.Model.State ...")
  return Cytof5.Model.State(Z=s.Z, delta=s.delta, alpha=s.alpha, v=s.v, W=s.W,
                            sig2=s.sig2, eta=s.eta, lam=s.lam, gam=s.gam,
                            y_imputed=s.y_imputed)
end

# FIXME: similarity_Z...
function convertConstants(c)
  println("Trying to convert reconstructed state to Cytof5.Model.Constants ...")
  return Cytof5.Model.Constants(alpha_prior=c.alpha_prior,
                                delta_prior=c.delta_prior, W_prior=c.W_prior,
                                eta_prior=c.eta_prior, sig2_prior=c.sig2_prior,
                                beta=c.beta, K=c.K, L=c.L,
                                sig2_range=c.sig2_range,
                                # there should be new stuff here ...
                                probFlip_Z=c.probFlip_Z, similarity_Z=x -> x)

end

# TODO: Write somrthing generic here.
function convertType(s, T)
  if T == Cytof5.Model.State
    return convertState(s)
  elseif T == Cytof5.Model.Constants
    return convertConstants(s)
  else
    println("Don't know how to convert to type $(T)")
  end
end

function post_process(path_to_output, thresh=0.9, min_presences=[0, .01, .03, .05])
  jld_filename = split(path_to_output, "/")[end]
  outdir = join(split(path_to_output, "/")[1:end-1], "/")
  # datapath = "data/reduced_cb.jld2"
  datapath = "$outdir/reduced_data/reduced_cb.jld2"
  println("datapath: $datapath")
  IMGDIR = "$outdir/img/"

  run(`mkdir -p $(IMGDIR)/dden/`)

  println("Loading Data ...")
  cbData = loadSingleObj(datapath)
  cbData = decompress_data(cbData)
  println("Loading Results ...")

  # FIXME: Not efficient!
  init_is_defined = "init" in keys(load(path_to_output))
  println("init_is_defined = $init_is_defined ...")

  if init_is_defined
    @load path_to_output out ll lastState c init dden
  else
    @load path_to_output out ll lastState c
  end

  I, K_MCMC = size(lastState.W)
  J = size(lastState.Z, 1)
  N = size.(lastState.y_imputed, 1)

  # Set png resolution settings
  s_png = 10
  fy(clus) = util.addCut(clus, s_png)
  fZ(Z) = util.addGridLines(Z, s_png)

  println("Making eta")
  open("$IMGDIR/eta.txt", "w") do file
    # TODO
    eta = util.getPosterior(:eta, out[1])
    eta0_mean = mean([ei[0] for ei in eta])
    eta1_mean = mean([ei[1] for ei in eta])
    eta_mean = Dict(0 => eta0_mean, 1 => eta1_mean)

    header = "i", "j", "z", "l", "eta"
    write(file, "$(join(header, "     "))\n")
    for i in 1:I
      for j in 1:J
        for z in 0:1
          for l in 1:c.L[z]
            line = (i, j, z, l, eta_mean[z][i, j, l])
            line = Printf.@sprintf("%d %5d %5d %5d     %.5f", line...);
            write(file, "$line\n")
          end
        end
      end
    end
  end

  println("Making eta_obs.txt")
  try
    open("$IMGDIR/eta_obs.txt", "w") do file
      header = "i", "j", "z", "l", "p"
      write(file, "$(join(header, "     "))\n")

      for i in 1:I
        for j in 1:J
          for z in 0:1
            for l in 1:c.L[z]
              idx_observed = util.idx_observed_ij(cbData, i, j)
              z_ij = out[1][end][:Z][j, out[1][end][:lam][i][idx_observed]]
              gam_ij = out[2][end][:gam][i][idx_observed, j]
              mij_sum = length(idx_observed)
              p = sum((gam_ij .== l) .& (z_ij .== z)) / mij_sum
              line = (i, j, z, l, p)
              line = Printf.@sprintf("%d %5d %5d %5d     %.5f", line...);
              write(file, "$line\n")
            end
          end
        end
      end
    end
  catch
    open("$IMGDIR/eta_obs.log", "w") do file
      write(file, "failed ... \n")
    end
  end

  # Plot loglikelihood
  util.plotPdf("$(IMGDIR)/ll_complete_history.pdf")
  util.plot(ll, ylab="log-likelihood", xlab="MCMC iteration", typ="l");
  util.abline(v=length(ll) - length(out[1]))
  util.devOff()

  util.plotPdf("$(IMGDIR)/ll_postburn.pdf")
  util.plot(ll[end-length(out[1]):end], ylab="log-likelihood", xlab="MCMC iteration", typ="l");
  util.devOff()

  function addGridLines(J::Int, K::Int, col="grey", lwd=1)
    util.abline(v=(1:K) .+ .5, h=(1:J) .+ .5, col=col, lwd=lwd);
  end

  # Plot Z
  Zpost = util.getPosterior(:Z, out[1])
  Zmean = mean(Zpost)

  println("Making Z ...")
  try
    util.plotPdf("$IMGDIR/Z_mean.pdf")
    util.myImage(Zmean, xlab="Features", ylab="Markers", addL=true, col=util.greys(11),
                 f=Z->addGridLines(J, K_MCMC));
    util.devOff()

    util.plotPdf("$IMGDIR/Z_mean_est_leftordered.pdf")
    util.myImage(Cytof5.Model.leftOrder((Zmean .> .5)*1),
                 xlab="Features", ylab="Markers", addL=true, col=util.greys(11),
                 f=Z->addGridLines(J, K_MCMC));
    util.devOff()
  catch
    println("Failed to make complete making Z ...")
  end

  println("Plot Z diffs")
  if init_is_defined
    util.plotPdf("$IMGDIR/Z_minus_init.pdf")
    for i in 1:length(Zpost)
      z = Zpost[i]
      util.myImage(z - init.Z, xlab="Features", ylab="Markers", addL=true, col=util.blueToRed(3),
                   zlim=[-1, 1], f=Z->addGridLines(J, K_MCMC), main="Z_$i - Z_0");
    end
    util.devOff()
  else
    println("Can't make Z-diffs because init isn't defined ...")
    run(`rm -f $IMGDIR/Z_minus_init.pdf`)
  end

  # Plot Posterior Density for i, j observed
  for i in 1:I
    for j in 1:J
      println("Plot posterior density for (i: $i, j: $j) observed ...")
      util.plotPdf("$(IMGDIR)/dden/dden_i$(i)_j$(j).pdf")
      dyij = util.density(filter(yij -> !isnan(yij), cbData[i][:, j]))
      dd_ij = hcat([dd[i, j] for dd in dden]...)
      pdyij_mean = R"rowMeans($dd_ij)" .+ 0
      pdyij_lower = R"apply($dd_ij, 1, quantile, .025)" .+ 0
      pdyij_upper = R"apply($dd_ij, 1, quantile, .975)" .+ 0
      h = maximum([dyij[:y] .+ 0; pdyij_upper])
      util.plot(c.y_grid, pdyij_mean, xlab="y", ylab="density", ylim=[0, h],
                main="i: $i, j: $j", col="blue", lwd=2, typ="l")
      util.colorBtwn(c.y_grid, pdyij_lower, pdyij_upper, from=-10, to=10,
                     col=util.rgba("blue", .3))
      util.lines(dyij, col="grey", lwd=2)
      util.devOff()
    end
  end

  for i in 1:I
    println("Separate graphs for yZ $(i)...")
    idx_best = R"estimate_ZWi_index($(out[1]), $i)"[1]
    Zi = out[1][idx_best][:Z]
    Wi = out[1][idx_best][:W][i,:]

    # Point Est for Wi and Zi. txt and pdf.
    open("$IMGDIR/W_$(i)_hat.txt", "w") do file
      write(file, "$(join(Wi, "\n"))\n")
    end

    open("$IMGDIR/W$(i)_hat_ordered_cumsum.txt", "w") do file
      ord = sortperm(Wi, rev=true)
      cs_wi_sorted = cumsum(Wi[ord])
      write(file, "num_features,k,wi,cumprop\n")
      for k in 1:K_MCMC
        write(file, "$(k),$(ord[k]),$(Wi[ord][k]),$(cs_wi_sorted[k])\n")
      end
    end

    open("$IMGDIR/Z$(i)_hat.txt", "w") do file
      for j in 1:J
        zj = join(Int.(Zi[j, :]), ",")
        write(file, "$(zj)\n")
      end
    end

    lami = out[1][idx_best][:lam][i]
    ord = sortperm(Wi, rev=true)
    lami = util.reorder_lami(ord, lami)

    for min_presence in min_presences
      common_celltypes = util.get_common_celltypes(Wi, thresh=min_presence,
                                                   filter_by_min_presence=true)
      println("common celltypes (min_presence > $min_presence): $common_celltypes")
      K_trunc = length(common_celltypes)

      util.plotPng("$IMGDIR/y_dat$(i)_only_minpresence$(min_presence).png", s=s_png)
      ord_yi = sortperm(lami)
      function f(y)
        util.addCut(lami, s_png)
        R"""
        axis(1, at=1:$J, label=1:$J, las=2, fg="grey", cex.axis=1.5)
        """
      end
      util.myImage(cbData[i][ord_yi[1 .<= lami[ord_yi] .<= K_trunc], :],
                   addL=true, f=f, xaxt="n",
                   zlim=[-4,4], col=util.blueToRed(9), na="black", xlab="markers",
                   var"cex.lab"=1.5, ylab="cells", var"cex.y.leg"=1.5);
      # util.myImage(cbData[i][sortperm(lami), :], addL=true, f=yi->util.addCut(lami),
      #              zlim=[-4,4], col=util.blueToRed(9), na="black", xlab="markers",
      #              ylab="cells");
      util.devOff()

      util.plotPdf("$IMGDIR/Z_hat$(i)_minpresence$(min_presence).pdf", w=5, h=10)
      util.myImage(Zi[:, common_celltypes], addL=false, ylab="markers", yaxt="n",
                   f=Z->addGridLines(J, K_trunc), xaxt="n", xlab="celltypes",
                   var"cex.lab"=1.5);

      perc = string.(round.(Wi[common_celltypes] * 100, digits=2), "%")
      R"""
      axis(3, at=1:$K_trunc, label=$(perc), las=1, fg="grey", cex.axis=1)
      axis(1, at=1:$K_trunc, label=$(common_celltypes), las=1,
           fg="grey", cex.axis=1)
      axis(2, at=1:$J, label=1:$J, las=2, fg="grey", cex.axis=1)
      """
      util.devOff()

      util.plotPdf("$IMGDIR/ZT_hat$(i)_minpresence$(min_presence).pdf")
      perc = string.(round.(Wi[common_celltypes] * 100, digits=2), "%")
      function fz(z)
        addGridLines(K_trunc, J)
        R"""
        axis(4, at=1:$K_trunc, label=$(perc), las=1, fg="grey", cex.axis=1.3)
        axis(2, at=1:$K_trunc, label=$(common_celltypes), las=1,
             fg="grey", cex.axis=1.3)
        axis(1, at=1:$J, label=1:$J, las=2, fg="grey", cex.axis=1.3)
        """
      end
      R"par(mar=c(5.1, 4, 2.1, 5.1))"
      util.myImage(Zi[:, common_celltypes]', addL=false, ylab="celltypes", yaxt="n",
                   f=fz, xaxt="n", xlab="markers", var"cex.lab"=1.5);
      R"par(mar=c(5, 4, 4, 2) + .1)"
      util.devOff()
    end
  end

  # Plot alpha
  alphaPost = util.getPosterior(:alpha, out[1])
  println("Making alpha ...")
  util.plotPdf("$(IMGDIR)/alpha.pdf")
  util.plotPost(alphaPost, ylab="density", xlab="alpha", main="");
  util.devOff()

  # Plot W
  Wpost = util.getPosterior(:W, out[1])
  W_mean = mean(Wpost)

  println("Making W...")
  util.plotPdf("$IMGDIR/W.pdf")

  cols = vcat([k .+ c.K .* (0:I-1) for k in 1:c.K]...)
  colnames = vec(["$i:$k" for i in 1:I, k in 1:c.K])
  W_post = hcat([hcat([w[i, :] for w in Wpost]...)' for i in 1:I]...)
  util.boxplot(W_post[:, cols], col=R"rep(2:$(I+1), $(c.K))", names=colnames, las=2, cex=.2);
  util.abline(v=(0:c.K) * I .+ .5, lty=2, col="grey")

  R"par(mfrow=c($I, 1), mar=c(5, 5.1, 0.5, 2.1))"
  for i in 1:I
    util.boxplot(hcat([w[i, :] for w in Wpost]...)', ylab="Posterior: W$i",
                 xlab=i<I ? "" : "Features",
                 col="steelblue", pch=20, cex=0);
  end
  R"par(mfrow=c(1, 1), mar=rcommon::mar.default())"
  util.devOff()

  # Missing Mechanism
  println("Making beta  ...")
  open("$IMGDIR/beta.txt", "w") do file
    for i in 1:I
      bi = join(c.beta[:, i], ", ")
      write(file, "for i=$i, beta = $bi \n")
    end
  end

  # y_imputed cache
  y_imputed = [ o[:y_imputed] for o in out[2] ]
  B = length(y_imputed)
  y_imputed_min = minimum(vcat([minimum.(yb) for yb in y_imputed]...))
  y_imputed_max = maximum(vcat([maximum.(yb) for yb in y_imputed]...))
  y_imputed_range = [y_imputed_min, y_imputed_max]

  # Plot missing mechanism
  println("Making prob miss  ...")
  util.plotPdf("$IMGDIR/prob_miss.pdf")
  R"par(mfrow=c($I, 1))"
  for i in 1:I
    util.plotProbMiss(c.beta, i, xlim=y_imputed_range)
  end
  R"par(mfrow=c(1,1))"
  util.devOff()

  # Get mus
  delta0Post = hcat([m[:delta][0] for m in out[1]]...)'
  delta1Post = hcat([m[:delta][1] for m in out[1]]...)'
  mus0Post = -cumsum(delta0Post, dims=2)
  mus1Post =  cumsum(delta1Post, dims=2)
  musPost = [ mus0Post mus1Post ]

  println("Making mus...")
  util.plotPdf("$IMGDIR/mus.pdf")
  util.boxplot(musPost, ylab="mu*", xlab="", xaxt="n", col="steelblue", pch=20, cex=0);
  util.abline(h=0, v=size(mus0Post, 2) + .5, col="grey30", lty=1);
  util.devOff()

  util.plotPdf("$IMGDIR/delta0.pdf")
  util.plotPosts(delta0Post);
  util.devOff()

  util.plotPdf("$IMGDIR/delta1.pdf")
  util.plotPosts(delta1Post);
  util.devOff()

  # Get sig2
  sig2Post = hcat(util.getPosterior(:sig2, out[1])...)'
  sig2Mean = mean(sig2Post, dims=1)
  sig2Sd = std(sig2Post, dims=1)

  println("Making sig2...")
  util.plotPdf("$IMGDIR/sig2.pdf")
  util.plotPosts(sig2Post);
  util.boxplot(sig2Post, ylab="sig2", xlab="sample",
               col="steelblue", pch=20, cex=0,
               var"cex.lab"=1.5, var"cex.axis"=1.5);
  util.devOff()

  println("Making v...")
  vPost = hcat(util.getPosterior(:v, out[1])...)'
  util.plotPdf("$IMGDIR/v.pdf")
  util.boxplot(vPost, ylab="v", xlab="", xaxt="n", col="steelblue", pch=20, cex=0);
  util.devOff()

  util.plotPdf("$IMGDIR/v_each.pdf")
  for k in 1:c.K
    util.plotPost(vPost[:, k], main="v_$k")
  end
  util.devOff()

  println("Making v_cumprod...")
  util.plotPdf("$IMGDIR/v_cumprod.pdf")
  util.boxplot(cumprod(vPost, dims=2),
               ylab="v_cumprod", xlab="", xaxt="n", col="steelblue", pch=20, cex=0);
  util.devOff()

  # Posterior of eps
  println("Making eps...")
  epsPost = hcat(util.getPosterior(:eps, out[1])...)'
  util.plotPdf("$IMGDIR/eps.pdf")
  util.plotPosts(epsPost);
  util.devOff()

  # Posterior of y_imputed
  println("Making ydatPost.pdf ...")
  util.plotPdf("$(IMGDIR)/ydatPost.pdf")
  R"par(mfrow=c(4,2))"
  for i in 1:I
    for j in 1:J
      print("\r i: $i, j: $j")
      numMissing = sum(cbData[i][:, j] .=== NaN)
      h = maximum(util.density(y_imputed[1][i][:, j])[:y])
      util.plot(util.density([cbData[i][:, j]; fill(-10, numMissing)], na=true, bw=.3),
                col="red", xlim=[-8,8], main="Y sample: $(i), marker: $(j),
                missing: $numMissing", bty="n", fg="grey", ylim=[0,h])
      for iter in 1:B
        yimp = y_imputed[iter]
        util.lines(util.density(yimp[i][:, j]), col=util.rgba("blue", .5))
      end
    end
  end
  R"par(mfrow=c(1,1))"
  println()
  util.devOff()

  idx_missing = [ findall(isnan.(cbData[i])) for i in 1:I ]
  num_missing_per_sample= 3
  idx_missing = [shuffle!(idx_missing[i])[1:num_missing_per_sample] for i in 1:I]
  idx_missing = [map(idx -> (i, idx[1], idx[2]), idx_missing[i]) for i in 1:I]
  idx_missing = vcat(idx_missing...)
  for idx in idx_missing
    i, n, j = idx
    y_inj = [y_imputed[b][i][n, j] for b in 1:B]
    util.plotPdf("$(IMGDIR)/y_trace_i$(i)_n$(n)_j$(j).pdf")
    util.hist(y_inj, col="blue", border="transparent",
              main="", xlab="y (i: 2, n: $n, j: $j)", ylab="counts")
    util.plot(y_inj, typ="l", xlab="Index", ylab="y (i: $i, n: $n, j: $j)", main="")
    util.devOff()
  end

  for i in 1:I
    idx_missing = findall(isnan.(cbData[i]))
    y_i = [y_imputed[b][i][idx_missing] for b in 1:B]
    util.plotPdf("$IMGDIR/y_imputed_hist_i$(i).pdf")
    R"par(mfrow=c($I, 1), mar=c(5, 5.1, 0.5, 2.1))"
    util.hist(mean(y_i),    xlim=y_imputed_range, xlab="means of imputed y for sample $i", ylab="counts", main="");
    util.hist(maximum(y_i), xlim=y_imputed_range, xlab="max of imputed y for sample $i", ylab="counts", main="");
    util.hist(minimum(y_i), xlim=y_imputed_range, xlab="min of imputed y for sample $i", ylab="counts", main="");
    R"par(mfrow=c(1,1), mar=mar.default())"
    util.devOff()
  end

  for i in 1:I
    util.plotPng("$IMGDIR/y_imputed$(i).png", s=s_png)
    util.yZ_inspect(out[1], i=i, lastState.y_imputed, zlim=[-4,4], using_zero_index=false,
                    thresh=thresh, col=util.blueToRed(9), fy=fy, fZ=fZ)
    util.devOff()
  end

  for i in 1:I
    util.plotPng("$IMGDIR/y_dat$(i).png", s=s_png)
    util.yZ_inspect(out[1], i=i, cbData, zlim=[-4,4], using_zero_index=false, na="black",
                    thresh=thresh, col=util.blueToRed(9), fy=fy, fZ=fZ)
    util.devOff()
  end

  println("Making yZ with Z mean...")
  if any(lastState.eps .== 0)
    for i in 1:I
      util.plotPng("$IMGDIR/y_dat$(i)_with_zmean.png", s=s_png)

      idx_best = R"estimate_ZWi_index($(out[1]), $i)"[1]
      Zi = out[1][idx_best][:Z]
      Wi = out[1][idx_best][:W][i,:]
      lami = out[1][idx_best][:lam][i]

      open("$IMGDIR/lam$(i)_best.txt", "w") do file
        for lam_in in lami
          write(file, "$(lam_in)\n")
        end
      end

      S = [findall(lami .== k) for k in 1:c.K]
      Zi_bar = mean([[mean(o[:Z][j, o[:lam][i][S[k]]]) for j in 1:J, k in 1:c.K] for o in out[1]])

      util.yZ(cbData[i], Zi_bar, Wi, lami, zlim=[-4,4], thresh=thresh, col=util.blueToRed(9),
              na="black", using_zero_index=false, col_Z=R"grey(seq(1, 0, len=11))",
              colorbar_Z=true, cex_z_leg=0.001, fy=fy, fZ=fZ)

      util.devOff()
    end
  else
    # Get lam
    println("Making lam...")
    lamPost = util.getPosterior(:lam, out[1])
    unique(lamPost)

    prop_lam0 = hcat([[mean(lam[i] .== 0) for lam in lamPost] for i in 1:I]...)
    sum_lam0 = hcat([[sum(lam[i] .== 0) for lam in lamPost] for i in 1:I]...)
    println("prop0 dim: $(size(prop_lam0))")
    util.plotPdf("$IMGDIR/lam0.pdf")
    util.boxplot(prop_lam0, ylab="P(lambda_i = 0 | data)",
                 xlab="sample", col="steelblue", pch=20, cex=0);
    util.devOff()

    util.plotPdf("$IMGDIR/lam0_counts.pdf")
    util.boxplot(sum_lam0, ylab="sum(lambda_i = 0 | data)",
                 xlab="sample", col="steelblue", pch=20, cex=0);
    util.devOff()

    # Plot lambda = 0 group
    println("Plot noisy group")
    for i in 1:I
      idx_best = R"estimate_ZWi_index($(out[1]), $i)"[1]
      lami = out[1][idx_best][:lam][i]

      open("$IMGDIR/lam$(i)_best.txt", "w") do file
        for lam_in in lami
          write(file, "$(lam_in)\n")
        end
      end

      idx0 = findall(lami .== 0)
      if length(idx0) > 0
        println("making y_dat$(i)_lam0.png")
        util.plotPng("$IMGDIR/y_dat$(i)_lam0.png", s=s_png)
        util.myImage(cbData[i][idx0, :], xlab="markers", ylab="cells",
                     na="black", col=util.blueToRed(9), addL=true, zlim=[-4, 4], nticks=11)
        util.devOff()

        util.plotPng("$IMGDIR/y_imputed$(i)_lam0.png", s=s_png)
        util.myImage(lastState.y_imputed[i][idx0, :], xlab="markers", ylab="cells",
                     na="black", col=util.blueToRed(9), addL=true, zlim=[-4, 4])
        util.devOff()
      else
        println("length(lam0) == 0. Not making y_dat$(i)_lam0.png!")
      end
    end
  end

  for i in 1:I
    println("Printing number of predominant celltypes for i = $i ...")
    idx_best = R"estimate_ZWi_index($(out[1]), $i)"[1]
    Wi = out[1][idx_best][:W][i,:]

    open("$IMGDIR/K_P$(round(Int, thresh * 100))_i$(i).txt", "w") do file
      cswi = cumsum(sort(Wi, rev=true))
      K_TOP = findfirst(cswi .> thresh)
      write(file, "K_TOP: $K_TOP \n")
    end
  end

  if init_is_defined
    println("Making y_dat init ...")
    for i in 1:I
      util.plotPng("$IMGDIR/y_dat$(i)_init.png", s=s_png)
      Zi = init.Z
      Wi = init.W[i, :]
      lami = init.lam[i]

      util.yZ(cbData[i], Zi, Wi, lami, zlim=[-4,4], thresh=thresh, col=util.blueToRed(9),
              na="black", using_zero_index=false, col_Z=R"grey(seq(1, 0, len=11))",
              colorbar_Z=true, cex_z_leg=0.001, fy=fy, fZ=fZ)
      util.devOff()
    end
  else
    println("Skipping this plot, init is not defined ...")
  end
end


text:
#gtk_ext.jl: GTK extensions (features missing from Gtk.jl library)
#-------------------------------------------------------------------------------

#GdkAtom "pointers":
const GTK_SELECTION_CLIPBOARD = Ptr{Nothing}(69)

#==Functions
===============================================================================#
function window_close(window::Gtk.GtkWindow)
	ccall((:gtk_window_close,Gtk.libgtk),Nothing,(Ptr{Gtk.GObject},),window)
	return
end

function clipboard_set_text(clipbrd::Ptr{Nothing}, v::String, _len::Int)
	ccall((:gtk_clipboard_set_text, Gtk.libgtk), Nothing, (Ptr{Gtk.GObject}, Ptr{UInt8}, Cint), clipbrd, v, _len)
	return clipbrd
end

#"ext": for customized implementations
function GtkMessageDialogLeaf_ext(msg::String, buttons::Integer=GtkButtonsType.OK,
		msgtype::Integer=GtkMessageType.INFO, parent = Gtk.GtkNullContainer())
	w = Gtk.GtkMessageDialogLeaf(ccall((:gtk_message_dialog_new, Gtk.libgtk), Ptr{Gtk.GObject},
		(Ptr{Gtk.GObject}, Cint, Cint, Cint, Ptr{UInt8}),
		parent, GtkDialogFlags.DESTROY_WITH_PARENT, msgtype, buttons, C_NULL)
	)
	set_gtk_property!(w, :text, msg)
	return w
end
function dialog_confirm(msg::String, parent = Gtk.GtkNullContainer())
	w = GtkMessageDialogLeaf_ext(msg, GtkButtonsType.YES_NO, GtkMessageType.WARNING, parent)
	answer = run(w)
	Gtk.destroy(w)
	return (GtkResponseType.YES == answer)
end
function dialog_inputbox(msg::String, entry_default::String, parent = Gtk.GtkNullContainer())
	w = GtkMessageDialogLeaf_ext(msg, GtkButtonsType.OK_CANCEL, GtkMessageType.INFO, parent)
	box = Gtk.content_area(w)
	entry = Gtk.GtkEntry(; text = entry_default)
	push!(box, entry)
	Gtk.showall(w)
	cancel = (run(w) != GtkResponseType.OK)
	_value = get_gtk_property(entry, :text, String)
	Gtk.destroy(w)
	if cancel; return nothing; end
	return _value
end

#Last line

ext: jl


text:
"""
    log_hypergeometric(a, b, x)

Compute the logarithm of the hypergeometric function
``{}_{2}F_{1}(-a, -b; -(a + b - 1); -exp(x))`` where `a` and `b` are positive
integers, and `a >= b`.
"""
function log_hypergeometric(
  a::F,
  b::F,
  x::F
)::F where {
  F <: AbstractFloat
}
  if b > 1
    # define H_{b} = y_{b} / y_{b - 1}, i. e. y_{b} = H_{b} * y_{b - 1}
    # y_{0} = 1, y_{1} = 1 + (a * z) / (a + 1 - k)
    tmp = log1pexp(x)

    # v = [log(H), log(y), compensation_error]
    v = [tmp, tmp, F(0)]

    n = F(1)
    while n < b
      n += 1
      kbn_2F1_iteration!(v, n, a, x, F(1))
    end

    # finally add the compensation error
    v[2] + v[3]
  elseif b == 1
    # 2F1(-a, -1; -a, -z) = 1 + z
    log1pexp(x)
  else
    0
  end
end

"""
    log_hypergeometric_joint(a, b, x)

Compute the logarithm of the three hypergeometric functions
``{}_{2}F_{1}(-a, -b; -(a + b - 1); -exp(x))``,
``{}_{2}F_{1}(-(a - 1), -(b - 1); -(a + b); -exp(x))``, and
``{}_{2}F_{1}(-(a - 2), -(b - 2); -(a + b + 1); -exp(x))`` at the same time.
`a` and `b` are positive integers with `a >= b`. We assume that `a` and `b` are
both greater than 1.
"""
function log_hypergeometric_joint(
  a::F,
  b::F,
  x::F
)::Tuple{F, F, F} where {
  F <: AbstractFloat
}
  # define H_{b} = y_{b} / y_{b - 1}, i. e. y_{b} = H_{b} * y_{b - 1}
  # y_{0} = 1, y_{1} = 1 + (a * z) / (a + 1 - k)

  # k = 1
  # 2F1(-a, -b; -(a + b - 1); -z)
  tmp = log1pexp(x)
  v1 = [tmp, tmp, F(0)]

  # k = 0
  # 2F1(-(a - 1), -(b - 1); -(a + b); -z)
  tmp = log1pexp(log((a - 1) / a) + x)
  v2 = [tmp, tmp, F(0)]

  if b > 2
    # k = -1
    # 2F1(-(a - 2), -(b - 2); -(a + b + 1); -z)
    tmp = log1pexp(log((a - 2) / a) + x)
    v3 = [tmp, tmp, F(0)]

    n = F(1)
    while n < b - 2
      n += 1
      kbn_2F1_iteration!(v1, n, a, x, F(1))
      kbn_2F1_iteration!(v2, n, a - 1, x, F(0))
      kbn_2F1_iteration!(v3, n, a - 2, x, F(-1))
    end

    # We still need few iterations
    kbn_2F1_iteration!(v1, b - 1, a, x, F(1))
    kbn_2F1_iteration!(v2, b - 1, a - 1, x, F(0))

    kbn_2F1_iteration!(v1, b, a, x, F(1))

    # add compensation errors and return the final values
    v1[2] + v1[3], v2[2] + v2[3], v3[2] + v3[3]
  else
    # log(2F1(-a, -2; -a; -z)) = v1[2] + H
    # log(2F1(-(a - 1), -1; -(a + 2); -z)) = v2[2]
    # log(2F1(-(a - 2), 0; -(a + 3); -z)) = 0
    kbn_2F1_iteration!(v1, F(2), a, x, F(1))
    v1[2], v2[2], 0
  end
end

"""
    kbn_2F1_iteration!(v, n, a, x, k)

Basic Kahan iteration operations for computing
``{}_{2}F_{1}(-a, -b; -(a + b - k); -exp(x))``.
"""
@inline function kbn_2F1_iteration!(
  v::Vector{F},
  n::F,
  a::F,
  x::F,
  k::F
) where {
  F <: AbstractFloat
}
  # new value for log(y_{b} / y_{b - 1})
  v[1] = log1pexp(x + log((a - n + 1) / (a + n - k)) +
                  log1pexp(log((n - 1) * (n - 1 - k) /
                               ((a - n + 1) * (a + n - k - 1))) - v[1]))

  # running sum including error
  s = v[2] + v[1]

  # Neumaier modification to Kahan summation
  v[3] += if abs(v[2]) >= abs(v[1])
    (v[1] - s) + v[2]
  else
    (v[2] - s) + v[1]
  end

  v[2] = s

  nothing
end


text:
## %PREDICTMD_GENERATED_BY%

using PredictMDExtra
PredictMDExtra.import_all()

using PredictMD
PredictMD.import_all()

# PREDICTMD IF INCLUDE TEST STATEMENTS
logger = Base.CoreLogging.current_logger_for_env(Base.CoreLogging.Debug, Symbol(splitext(basename(something(@__FILE__, "nothing")))[1]), something(@__MODULE__, "nothing"))
if isnothing(logger)
    logger_stream = devnull
else
    logger_stream = logger.stream
end
# PREDICTMD ELSE
# PREDICTMD ENDIF INCLUDE TEST STATEMENTS

### Begin project-specific settings

DIRECTORY_CONTAINING_THIS_FILE = @__DIR__
PROJECT_DIRECTORY = dirname(
    joinpath(splitpath(DIRECTORY_CONTAINING_THIS_FILE)...)
    )
PROJECT_OUTPUT_DIRECTORY = joinpath(
    PROJECT_DIRECTORY,
    "output",
    )
mkpath(PROJECT_OUTPUT_DIRECTORY)
mkpath(joinpath(PROJECT_OUTPUT_DIRECTORY, "data"))
mkpath(joinpath(PROJECT_OUTPUT_DIRECTORY, "models"))
mkpath(joinpath(PROJECT_OUTPUT_DIRECTORY, "plots"))

# PREDICTMD IF INCLUDE TEST STATEMENTS
@debug("PROJECT_OUTPUT_DIRECTORY: ", PROJECT_OUTPUT_DIRECTORY,)
if PredictMD.is_travis_ci()
    PredictMD.cache_to_path!(
        ;
        from = ["cpu_examples", "breast_cancer_biopsy", "output",],
        to = [PROJECT_OUTPUT_DIRECTORY],
        )
end
# PREDICTMD ELSE
# PREDICTMD ENDIF INCLUDE TEST STATEMENTS

### End project-specific settings

### Begin SMOTE class-balancing code

Random.seed!(999)

trainingandtuning_features_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "trainingandtuning_features_df.csv",
    )
trainingandtuning_labels_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "trainingandtuning_labels_df.csv",
    )
testing_features_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "testing_features_df.csv",
    )
testing_labels_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "testing_labels_df.csv",
    )
training_features_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "training_features_df.csv",
    )
training_labels_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "training_labels_df.csv",
    )
tuning_features_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "tuning_features_df.csv",
    )
tuning_labels_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "tuning_labels_df.csv",
    )
trainingandtuning_features_df = DataFrames.DataFrame(
    FileIO.load(
        trainingandtuning_features_df_filename;
        type_detect_rows = 100,
        )
    )
trainingandtuning_labels_df = DataFrames.DataFrame(
    FileIO.load(
        trainingandtuning_labels_df_filename;
        type_detect_rows = 100,
        )
    )
testing_features_df = DataFrames.DataFrame(
    FileIO.load(
        testing_features_df_filename;
        type_detect_rows = 100,
        )
    )
testing_labels_df = DataFrames.DataFrame(
    FileIO.load(
        testing_labels_df_filename;
        type_detect_rows = 100,
        )
    )
training_features_df = DataFrames.DataFrame(
    FileIO.load(
        training_features_df_filename;
        type_detect_rows = 100,
        )
    )
training_labels_df = DataFrames.DataFrame(
    FileIO.load(
        training_labels_df_filename;
        type_detect_rows = 100,
        )
    )
tuning_features_df = DataFrames.DataFrame(
    FileIO.load(
        tuning_features_df_filename;
        type_detect_rows = 100,
        )
    )
tuning_labels_df = DataFrames.DataFrame(
    FileIO.load(
        tuning_labels_df_filename;
        type_detect_rows = 100,
        )
    )

categorical_feature_names_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "categorical_feature_names.jld2",
    )
continuous_feature_names_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "continuous_feature_names.jld2",
    )
categorical_feature_names = FileIO.load(
    categorical_feature_names_filename,
    "categorical_feature_names",
    )
continuous_feature_names = FileIO.load(
    continuous_feature_names_filename,
    "continuous_feature_names",
    )
feature_names = vcat(categorical_feature_names, continuous_feature_names)

single_label_name = :Class
negative_class = "benign"
positive_class = "malignant"
single_label_levels = [negative_class, positive_class]

categorical_label_names = Symbol[single_label_name]
continuous_label_names = Symbol[]
label_names = vcat(categorical_label_names, continuous_label_names)

DataFrames.describe(logger_stream, training_labels_df[single_label_name])
show(logger_stream, StatsBase.countmap(training_labels_df[single_label_name]))

majorityclass = "benign"
minorityclass = "malignant"

(smoted_training_features_df, smoted_training_labels_df,) = PredictMD.smote(
    training_features_df,
    training_labels_df,
    feature_names,
    single_label_name;
    majorityclass = majorityclass,
    minorityclass = minorityclass,
    pct_over = 100,
    minority_to_majority_ratio = 1.0,
    k = 5,
    )

PredictMD.check_column_types(
    smoted_training_features_df;
    categorical_feature_names = categorical_feature_names,
    continuous_feature_names = continuous_feature_names,
    categorical_label_names = categorical_label_names,
    continuous_label_names = continuous_label_names,
    )
PredictMD.check_column_types(
    smoted_training_labels_df;
    categorical_feature_names = categorical_feature_names,
    continuous_feature_names = continuous_feature_names,
    categorical_label_names = categorical_label_names,
    continuous_label_names = continuous_label_names,
    )

# PREDICTMD IF INCLUDE TEST STATEMENTS
Test.@test PredictMD.check_no_constant_columns(smoted_training_features_df)
Test.@test PredictMD.check_no_constant_columns(smoted_training_labels_df)
# PREDICTMD ELSE
# PREDICTMD ENDIF INCLUDE TEST STATEMENTS

DataFrames.describe(logger_stream, smoted_training_labels_df[single_label_name])
show(logger_stream, StatsBase.countmap(smoted_training_labels_df[single_label_name]))

smoted_training_features_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "smoted_training_features_df.csv",
    )
smoted_training_labels_df_filename = joinpath(
    PROJECT_OUTPUT_DIRECTORY,
    "data",
    "smoted_training_labels_df.csv",
    )
FileIO.save(smoted_training_features_df_filename, smoted_training_features_df)
FileIO.save(smoted_training_labels_df_filename, smoted_training_labels_df)

### End SMOTE class-balancing code

# PREDICTMD IF INCLUDE TEST STATEMENTS
if PredictMD.is_travis_ci()
    PredictMD.path_to_cache!(
        ;
        to = ["cpu_examples", "breast_cancer_biopsy", "output",],
        from = [PROJECT_OUTPUT_DIRECTORY],
        )
end
# PREDICTMD ELSE
# PREDICTMD ENDIF INCLUDE TEST STATEMENTS

## %PREDICTMD_GENERATED_BY%


text:
using AlgebraPDF
using StaticArrays
using Test
using SpecialFunctions
using LinearAlgebra
using Measurements

include("testparameters.jl")
include("testfuncions.jl")
include("testsum.jl")
include("testpdf.jl")
include("testdensities.jl")
include("testconvolution.jl")
include("testplotting.jl")
include("testmultidims.jl")
include("testfit.jl")
include("testgenerate.jl")


text:
"""
Read a table from another data structure that supports indexing columns by
symbols, such as a DataFrame. Requires a `Val{Names}` where `Names` is a
tuple of Symbols.

The reverse transformation (to DataFrames) can be acheived by:
   DataFrame(collect(table.data), collect(names(table)))
"""
@generated function readtable{Names}(::Type{Val{Names}}, x)
    exprs = ntuple(i->:(x[$(Expr(:quote,Names[i]))]), length(Names))
    expr = Expr(:tuple, exprs...)
    return :( Table{Names}($expr) )
end

"""
Read a table from a file in DLM (delimited) or CSV (comma-sereperated values,
default) text-based formats, using Julia's inbuilt readdlm. Set `multiplespaces = true`
instead of `delim = ' '` for files with arbitrarily many spaces between fields.
"""
function readtable{Names,Types}(::Type{Table{Names,Types}}, file::Union{AbstractString,IO}; kwargs...)
    table = Table{Names,Types}()
    readtable!(table, file; kwargs...)
    return table
end

function readtable!(table::Table, filename::AbstractString; kwargs...)
    open(filename) do fio
        readtable!(table, fio; kwargs...)
    end
end

"""
Append data to a table from a file in DLM (delimited) or CSV (comma-sereperated
values, default) text-based formats, using Julia's inbuilt readdlm. Set `multiplespaces = true`
instead of `delim = ' '` for files with arbitrarily many spaces between fields.
"""
function readtable!{Names, Types}(table::Table{Names, Types}, file::IO; delim::Char = ',', eol::Char = '\n', header::Bool = false, multiplespaces = false, kwargs...)
    if multiplespaces == true
        delim = Base.DataFmt.invalid_dlm(Char)
    end

    # Read data from file...
    local headerdata, data
    if header == true
        (data, headerdata) = readdlm(file, delim, Any, eol; header = true, kwargs...)
    else
        data = readdlm(file, delim, Any, eol; header = false, kwargs...)
    end

    # Read the header, verify columns and get corresponding indices
    local indices
    if header == true
        # Reoreded or extra indices are possible
        colnames = map(string, Names)

        indices = Vector{Int}(length(Names))
        for i = 1:length(colnames)
            found = 0
            for j = 1:length(headerdata)
                if colnames[i] == headerdata[j]
                    indices[i] = j
                    found += 1
                end
            end
            if found == 0
                error("Didn't find header string $(colnames[i])")
            elseif found > 1
                error("Found $found duplicates of header string $(colnames[i])")
            end
        end
    else
        indices = collect(1:length(Names))
    end

    # Convert to a table
    for i = 1:size(data,1)
        push!(table, ntuple(j->convert(eltype(Types.parameters[j]), data[i,indices[j]]), length(Names)))
    end
end

function writetable(filename::AbstractString, table::Table; kwargs...)
    fio = open(filename, false, true, true, true, false)
    writetable(fio, table; kwargs...)
    close(fio)
end

"""
Write a table to disk using delimeted text format. Optional arguments include:

  header = false
  delim = ','
  eol = '\n'
  null = "NA"
  string_delim = '\"'
  string_delim_right = string_delim (for guillements, chevrons, Gänsefüßchen...)
  char_delim = '\''
  char_delim_right = char_delim
"""
function writetable{Names,Types}(fileio::IO, table::Table{Names,Types}; header::Bool = false, delim = ',', eol = '\n', string_delim = '\"', string_delim_right = string_delim, char_delim = '\'', char_delim_right = char_delim, null = "NA")
    if header == true
        for j = 1:ncol(table)
            write(fileio, Names[j])
            if j == ncol(table)
                write(fileio, eol)
            else
                write(fileio, delim)
            end
        end
    end

    isnullable = fill(false, ncol(table))
    isstring = fill(false, ncol(table))
    ischar = fill(false, ncol(table))
    for i = 1:length(Names)
        typ = eltype(Types.parameters[i])
        if typ <: Nullable
            isnullable[i] = true
            typ = eltype(typ)
        end
        isstring[i] = typ <: AbstractString
        ischar[i] = typ <: Char
    end

    for i = 1:nrow(table)
        for j = 1:ncol(table)
            if isnullable[j]
                if isnull(table.data[j][i])
                    write(fileio, null)

                    if j == ncol(table)
                        write(fileio, eol)
                    else
                        write(fileio, delim)
                    end
                else
                    if isstring[j]
                        write(fileio, string_delim)
                    elseif ischar[j]
                        write(fileio, char_delim)
                    end
                    write(fileio, string(get(table.data[j][i])))
                    if isstring[j]
                        write(fileio, string_delim_right)
                    elseif ischar[j]
                        write(fileio, char_delim_right)
                    end

                    if j == ncol(table)
                        write(fileio, eol)
                    else
                        write(fileio, delim)
                    end
                end
            else
                if isstring[j]
                    write(fileio, string_delim)
                elseif ischar[j]
                    write(fileio, char_delim)
                end
                write(fileio, string(table.data[j][i]))
                if isstring[j]
                    write(fileio, string_delim_right)
                elseif ischar[j]
                    write(fileio, char_delim_right)
                end

                if j == ncol(table)
                    write(fileio, eol)
                else
                    write(fileio, delim)
                end
            end
        end
    end
end


text:
using AbstractTrees
using YAML

export build_tree, ComputeNode, ComputeGraph, get_commands, params
# In this file we define a tree data structure which can be used for optimisation passes
# over contraction commands

"""Generic tree node data structure"""
mutable struct ComputeNode{T}
    op::Union{Nothing, T}
    children::Vector{ComputeNode}
    parent::ComputeNode
    # Root constructor
    ComputeNode{T}(data) where T = new{T}(data, ComputeNode[])
    ComputeNode{T}() where T = new{T}(nothing, ComputeNode[])
    # Child node constructor
    ComputeNode{T}(data, parent::ComputeNode{U}) where {T, U} = new{T}(data, ComputeNode[], parent)
end
ComputeNode(op) = ComputeNode{typeof(op)}(op)

"""Represent a compute graph with root node and initial tensors"""
struct ComputeGraph
    root::ComputeNode
    tensors::Dict{Symbol, AbstractArray}
end

"""
    get_commands(ct::ComputeNode, type::Type=Any; by=nothing, iterf=PostOrderDFS)

Utility function that retrieves a list of commands if given type and sorted according to the
provided criteria. By default they are returned in the order returned by depth first traversal
which returns leaves before parents.
"""
function get_commands(cn::ComputeNode, type::Type=Any; by=nothing, iterf=PostOrderDFS)
    cmds = map(x -> x.op, filter(x -> x.op isa type, collect(iterf(cn))))
    if by !== nothing
        sort!(cmds, by=by)
    end
    cmds
end

"""Implement for ComputeGraph also"""
get_commands(cg::ComputeGraph, args...; kwargs...) = get_commands(cg.root, args...; kwargs...)

"""
    AbstractTrees.children(node::ComputeNode)

Implement children function from AbstractTrees package
"""
function AbstractTrees.children(node::ComputeNode)
    Tuple(node.children)
end

"""
    params(node::ComputeNode, optype::Type=Any)

Compile all parameters from this node and descendents with optional op type
qualifier. This can be used to return only output or view parameters with

params(node, OutputCommand)

or

params(node, ViewCommand)
"""
function params(node::ComputeNode, optype::Type=Any)
    local_params = node.op isa optype ? params(node.op) : Dict{Symbol, Int}()
    merge!(local_params, params.(node.children, [optype])...)
end

params(cg::ComputeGraph, args...) = params(cg.root, args...)
output(node::ComputeNode) = output(node.op)
output(cg::ComputeGraph) = output(cg.root)

###########################################################################
# Functions and data structures for trees of contraction commands
###########################################################################

"""
    build_tree(cmds::Vector{<: AbstractCommand})

Function to construct a tree from a list of commands
"""
function build_tree(cmds::Vector{<: AbstractCommand})
    nodes = Dict{Symbol, ComputeNode}()

    for op in cmds
        node = ComputeNode(op)
        for input in inputs(op)
            if haskey(nodes, input)
                push!(node.children, nodes[input])
                nodes[input].parent = node
            end
        end
        nodes[output(op)] = node
    end
    parentless = collect(keys(filter(x -> !isdefined(x[2], :parent), nodes)))
    @assert length(parentless) == 1 "Only root node should have no parent"
    root = parentless[1]
    nodes[root]
end


text:
"""
# module laminart

- Julia version: 1.4
- Author: niallcullinane
- Date: 2020-06-07

# Examples

```jldoctest
julia>
```
"""

module LaminartGPU

include("./LamKernels.jl")

using NNlib, ImageFiltering, Images, OffsetArrays, CUDA
# , MEngProject.LamKernels

export I_u, fun_v_C, fun_equ

mutable struct LamFunction{T <: AbstractArray} <: Function
	x::T
	m::T
	s::T

	x_lgn::T
	C::T
	H_z::T

	dy_temp::T
	dm_temp::T
	dz_temp::T
	ds_temp::T
	dv_temp::T
	H_z_temp::T
	V_temp_1::T
	V_temp_2::T
	Q_temp::T
	P_temp::T
end

function (ff::LamFunction)(du, u, p, t)

       @inbounds begin

		@. ff.x = @view u[:, :, 1:p.K,:]
        y = @view u[:, :, p.K+1:2*p.K,:]
        @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
        z = @view u[:, :, 3*p.K+1:4*p.K,:]
		@. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

		v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
		v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

		dx = @view du[:, :, 1:p.K,:]
        dy = @view du[:, :, p.K+1:2*p.K,:]
        dm = @view du[:, :, 2*p.K+1:3*p.K,:]
        dz = @view du[:, :, 3*p.K+1:4*p.K,:]
        ds = @view du[:, :, 4*p.K+1:5*p.K,:]

		dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
		dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

		fun_x_lgn!(ff.x_lgn, ff.x, p)
        fun_v_C!(ff.C, v_p, v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
        fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

        fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.dv_temp, p)
        fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
        fun_dx_v1!(dx, ff.x, ff.C, z, p.x_V2, p)
        fun_dy!(dy, y, ff.C, ff.x, ff.m, ff.dy_temp, p)
        fun_dm!(dm, ff.m, ff.x, ff.dm_temp, p)
        fun_dz!(dz, z, y, ff.H_z, ff.s, ff.dz_temp, p)
        fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

    end
    return nothing

end

mutable struct LamFunction_allStruct <: Function
	x::AbstractArray
	y::AbstractArray
	m::AbstractArray
	z::AbstractArray
	s::AbstractArray
	v_p::AbstractArray
	v_m::AbstractArray

	dx::AbstractArray
	dy::AbstractArray
	dm::AbstractArray
	dz::AbstractArray
	ds::AbstractArray
	dv_p::AbstractArray
	dv_m::AbstractArray

	x_lgn::AbstractArray
	C::AbstractArray
	H_z::AbstractArray

	dy_temp::AbstractArray
	dm_temp::AbstractArray
	dz_temp::AbstractArray
	ds_temp::AbstractArray
	dv_temp::AbstractArray
	H_z_temp::AbstractArray
	V_temp_1::AbstractArray
	V_temp_2::AbstractArray
	Q_temp::AbstractArray
	P_temp::AbstractArray
end

function (ff::LamFunction_allStruct)(du, u, p, t)

       @inbounds begin

		@. ff.x = @view u[:, :, 1:p.K,:]
        ff.y = @view u[:, :, p.K+1:2*p.K,:]
        @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
        ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
		@. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

		ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
		ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

		ff.dx = @view du[:, :, 1:p.K,:]
        ff.dy = @view du[:, :, p.K+1:2*p.K,:]
        ff.dm = @view du[:, :, 2*p.K+1:3*p.K,:]
        ff.dz = @view du[:, :, 3*p.K+1:4*p.K,:]
        ff.ds = @view du[:, :, 4*p.K+1:5*p.K,:]

		ff.dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
		ff.dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

		fun_x_lgn!(ff.x_lgn, ff.x, p)
        fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
        fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

        fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, ff.dv_temp, p)
        fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
        fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
        fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
        fun_dm!(ff.dm, ff.m, ff.x, ff.dm_temp, p)
        fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, ff.dz_temp, p)
        fun_ds!(ff.ds, ff.s, ff.H_z, ff.ds_temp, p)

    end
    return nothing

end

mutable struct LamFunction_gpu_reuse{T <: AbstractArray} <: Function

	x_lgn::T
	C::T
	H_z::T

	tmp_a::T
	tmp_b::T

	tmp_A::T
	tmp_B::T
	tmp_C::T
end

function (ff::LamFunction_gpu_reuse)(du, u, p, t)

       @inbounds begin

        y = @view u[:, :, p.K+1:2*p.K,:]
		z = @view u[:, :, 3*p.K+1:4*p.K,:]

		v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
		v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

		dx = @view du[:, :, 1:p.K,:]
        dy = @view du[:, :, p.K+1:2*p.K,:]
        dm = @view du[:, :, 2*p.K+1:3*p.K,:]
		dz = @view du[:, :, 3*p.K+1:4*p.K,:]
        ds = @view du[:, :, 4*p.K+1:5*p.K,:]

		dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
		dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

        fun_v_C!(ff.C, v_p, v_m, ff.tmp_a, ff.tmp_b, ff.tmp_A, ff.tmp_B, p) #a,b,C,D: buffers
        fun_H_z!(ff.H_z, z, ff.tmp_A, p) #A=buffer

		@. ff.tmp_A = @view u[:, :, 1:p.K,:] #x
        @. ff.tmp_B = @view u[:, :, 2*p.K+1:3*p.K,:] #m
		fun_x_lgn!(ff.x_lgn, ff.tmp_A, p)

        fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.tmp_a, p) #a=buffer
        fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.tmp_a, p) #a=buff

		fun_dx_v1!(dx, ff.tmp_A, ff.C, z, p.x_V2, p) #A:x
        fun_dy!(dy, y, ff.C, ff.tmp_A, ff.tmp_B, ff.tmp_C, p) # A:x, B:m, C:buffer
        fun_dm!(dm, ff.tmp_B, ff.tmp_A, ff.tmp_C, p) # A:x, B:m, C:buffer

		@. ff.tmp_A = @view u[:, :, 4*p.K+1:5*p.K,:] #s
		fun_dz!(dz, z, y, ff.H_z, ff.tmp_A, ff.tmp_B, p) #A=s, B=buffer
        fun_ds!(ds, ff.tmp_A, ff.H_z, ff.tmp_B, p) #A=s, B=buffer

    end
    return nothing

end

mutable struct LamFunction_equ{T <: AbstractArray} <: Function
	x::T
	m::T
	s::T

	x_lgn::T
	C::T
	H_z::T

	dy_temp::T
	dm_temp::T
	dz_temp::T
	ds_temp::T
	dv_temp::T
	H_z_temp::T
	V_temp_1::T
	V_temp_2::T
	Q_temp::T
	P_temp::T
end

function (ff::LamFunction_equ)(du, u, p, t)

       @inbounds begin

		@. ff.x = @view u[:, :, 1:p.K,:]
        y = @view u[:, :, p.K+1:2*p.K,:]
        @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
        z = @view u[:, :, 3*p.K+1:4*p.K,:]
		@. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

		v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
		v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

		dx = @view du[:, :, 1:p.K,:]
        dy = @view du[:, :, p.K+1:2*p.K,:]
        dm = @view du[:, :, 2*p.K+1:3*p.K,:]
        dz = @view du[:, :, 3*p.K+1:4*p.K,:]
        ds = @view du[:, :, 4*p.K+1:5*p.K,:]

		dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
		dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

		fun_x_lgn!(ff.x_lgn, ff.x, p)
        fun_v_C!(ff.C, v_p, v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
        fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

        fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.dv_temp, p)
        fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dx_v1!(dx, ff.x, ff.C, z, p.x_V2, p)
#         fun_dy!(dy, y, ff.C, ff.x, ff.m, ff.dy_temp, p)
		fun_x_equ!(ff.x, ff.C, z, ff.dy_temp, p.x_V2, p)
		fun_y_equ!(y, ff.C, ff.x, ff.m, ff.dy_temp, p)
        fun_dm!(dm, ff.m, ff.x, ff.dm_temp, p)
        fun_dz!(dz, z, y, ff.H_z, ff.s, ff.dz_temp, p)
        fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

    end
    return nothing

end

mutable struct LamFunction_all_struct_reuse_1{T <: AbstractArray} <: Function

	x_lgn::T
	C::T
	H_z::T

	y
	z
	v_p
	v_m
	dx
	dy
	dm
	dz
	ds
	dv_p
	dv_m

	tmp_a::T
	tmp_b::T

	tmp_A::T
	tmp_B::T
	tmp_C::T
end

function (ff::LamFunction_all_struct_reuse_1)(du, u, p, t)

       @inbounds begin

        ff.y = @view u[:, :, p.K+1:2*p.K,:]
		ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]

		ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
		ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

		ff.dx = @view du[:, :, 1:p.K,:]
        ff.dy = @view du[:, :, p.K+1:2*p.K,:]
        ff.dm = @view du[:, :, 2*p.K+1:3*p.K,:]
		ff.dz = @view du[:, :, 3*p.K+1:4*p.K,:]
        ff.ds = @view du[:, :, 4*p.K+1:5*p.K,:]

		ff.dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
		ff.dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

        fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.tmp_a, ff.tmp_b, ff.tmp_A, ff.tmp_B, p) #a,b,C,D: buffers
        fun_H_z!(ff.H_z, ff.z, ff.tmp_A, p) #A=buffer

		@. ff.tmp_A = @view u[:, :, 1:p.K,:] #x
        @. ff.tmp_B = @view u[:, :, 2*p.K+1:3*p.K,:] #m
		fun_x_lgn!(ff.x_lgn, ff.tmp_A, p)

        fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, ff.tmp_a, p) #a=buffer
        fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.tmp_a, p) #a=buff

		fun_dx_v1!(ff.dx, ff.tmp_A, ff.C, ff.z, p.x_V2, p) #A:x
        fun_dy!(ff.dy, ff.y, ff.C, ff.tmp_A, ff.tmp_B, ff.tmp_C, p) # A:x, B:m, C:buffer
        fun_dm!(ff.dm, ff.tmp_B, ff.tmp_A, ff.tmp_C, p) # A:x, B:m, C:buffer

		@. ff.tmp_A = @view u[:, :, 4*p.K+1:5*p.K,:] #s
		fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.tmp_A, ff.tmp_B, p) #A=s, B=buffer
        fun_ds!(ff.ds, ff.tmp_A, ff.H_z, ff.tmp_B, p) #A=s, B=buffer

    end
    return nothing

end

# mutable struct LamFunction_gpu30 <: Function
# 	x_lgn::AbstractArray
# 	C::AbstractArray
# 	H_z::AbstractArray

# 	dy_temp::AbstractArray
# 	dm_temp::AbstractArray
# 	dz_temp::AbstractArray
# 	ds_temp::AbstractArray
# 	dv_temp::AbstractArray
# 	H_z_temp::AbstractArray
# 	V_temp_1::AbstractArray
# 	V_temp_2::AbstractArray
# 	Q_temp::AbstractArray
# 	P_temp::AbstractArray
# end

# function (ff::LamFunction_gpu30)(du, u, p, t)

#        @inbounds begin

# 		ff.x = @view u[:, :, 1:p.K,:]
#         ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
# 		@. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

# 		ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		ff.dx = @view du[:, :, 1:p.K,:]
#         ff.dy = @view du[:, :, p.K+1:2*p.K,:]
#         ff.dm = @view du[:, :, 2*p.K+1:3*p.K,:]
#         ff.dz = @view du[:, :, 3*p.K+1:4*p.K,:]
#         ff.ds = @view du[:, :, 4*p.K+1:5*p.K,:]

# 		ff.dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
# 		ff.dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

# 		fun_x_lgn!(ff.x_lgn, ff.x, p)
#         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

#         fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
#         fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
#         fun_dm!(ff.dm, ff.m, ff.x, ff.dm_temp, p)
#         fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, ff.dz_temp, p)
#         fun_ds!(ff.ds, ff.s, ff.H_z, ff.ds_temp, p)

#     end
#     return nothing

# end

# mutable struct LamFunction_gpu14{T <: CuArray{Float32,4,Nothing}} <: Function
# 	x::T
# 	y::T
# 	m::T
# 	z::T
# 	s::T
# 	v_p::T
# 	v_m::T

# # 	dx::T
# # 	dy::T
# # 	dm::T
# # 	dz::T
# # 	ds::T
# # 	dv_p::T
# # 	dv_m::T

# 	x_lgn::T
# 	C::T
# 	H_z::T

# 	dy_temp::T
# 	dm_temp::T
# 	dz_temp::T
# 	ds_temp::T
# 	dv_temp::T
# 	H_z_temp::T
# 	V_temp_1::T
# 	V_temp_2::T
# 	Q_temp::T
# 	P_temp::T
# end

# function (ff::LamFunction_gpu14)(du, u, p, t)

#        @inbounds begin
# #         @. ff.x = @view u[:, :, 1:p.K,:]
# #         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
# #         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

# #        @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# # 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		x = @view u[:, :, 1:p.K,:]
#         y = @view u[:, :, p.K+1:2*p.K,:]
# #         m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         z = @view u[:, :, 3*p.K+1:4*p.K,:]
# #         s = @view u[:, :, 4*p.K+1:5*p.K,:]

#        v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		 dx = @view du[:, :, 1:p.K,:]
#         dy = @view du[:, :, p.K+1:2*p.K,:]
#         dm = @view du[:, :, 2*p.K+1:3*p.K,:]
#         dz = @view du[:, :, 3*p.K+1:4*p.K,:]
#         ds = @view du[:, :, 4*p.K+1:5*p.K,:]

#        dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
# 		dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

# #         fun_x_lgn!(ff.x_lgn, ff.x, p)
# #         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
# #         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

# #         fun_dv!(dv_p, ff.v_p, p.r, ff.x_lgn, ff.dv_temp, p)
# #         fun_dv!(dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
# #         fun_dx_v1!(dx, ff.x, ff.C, ff.z, p.x_V2, p)
# #         fun_dy!(dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
# #         fun_dm!(dm, ff.m, ff.x, ff.dm_temp, p)
# #         fun_dz!(dz, ff.z, ff.y, ff.H_z, ff.s, ff.dz_temp, p)
# #         fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

# # 		fun_x_lgn!(ff.x_lgn, x, p)
#         fun_v_C!(ff.C, v_p, v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

#         fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dx_v1!(dx, x, ff.C, z, p.x_V2, p)
#         fun_dy!(dy, y, ff.C, x, ff.m, ff.dy_temp, p)
#         fun_dm!(dm, ff.m, x, ff.dm_temp, p)
#         fun_dz!(dz, z, y, ff.H_z, ff.s, ff.dz_temp, p)
#         fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

# # 		@. du[:, :, 1:p.K,:] = ff.dx
# # 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# # 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# # 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# # 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

# #         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# # 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing

# end

# mutable struct LamFunction_gpu14{T <: CuArray{Float32,4,Nothing}} <: Function
# 	x::T
# 	y::T
# 	m::T
# 	z::T
# 	s::T
# 	v_p::T
# 	v_m::T

# # 	dx::T
# # 	dy::T
# # 	dm::T
# # 	dz::T
# # 	ds::T
# # 	dv_p::T
# # 	dv_m::T

# 	x_lgn::T
# 	C::T
# 	H_z::T

# 	dy_temp::T
# 	dm_temp::T
# 	dz_temp::T
# 	ds_temp::T
# 	dv_temp::T
# 	H_z_temp::T
# 	V_temp_1::T
# 	V_temp_2::T
# 	Q_temp::T
# 	P_temp::T
# end

# function (ff::LamFunction_gpu14)(du, u, p, t)

#        @inbounds begin
# #         @. ff.x = @view u[:, :, 1:p.K,:]
# #         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
# #         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

# #        @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# # 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		x = @view u[:, :, 1:p.K,:]
#         y = @view u[:, :, p.K+1:2*p.K,:]
# #         m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         z = @view u[:, :, 3*p.K+1:4*p.K,:]
# #         s = @view u[:, :, 4*p.K+1:5*p.K,:]

#        v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		 dx = @view du[:, :, 1:p.K,:]
#         dy = @view du[:, :, p.K+1:2*p.K,:]
#         dm = @view du[:, :, 2*p.K+1:3*p.K,:]
#         dz = @view du[:, :, 3*p.K+1:4*p.K,:]
#         ds = @view du[:, :, 4*p.K+1:5*p.K,:]

#        dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
# 		dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

# #         fun_x_lgn!(ff.x_lgn, ff.x, p)
# #         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
# #         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

# #         fun_dv!(dv_p, ff.v_p, p.r, ff.x_lgn, ff.dv_temp, p)
# #         fun_dv!(dv_m, ff.v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
# #         fun_dx_v1!(dx, ff.x, ff.C, ff.z, p.x_V2, p)
# #         fun_dy!(dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
# #         fun_dm!(dm, ff.m, ff.x, ff.dm_temp, p)
# #         fun_dz!(dz, ff.z, ff.y, ff.H_z, ff.s, ff.dz_temp, p)
# #         fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

# # 		fun_x_lgn!(ff.x_lgn, x, p)
#         fun_v_C!(ff.C, v_p, v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, z, ff.H_z_temp, p)

#         fun_dv!(dv_p, v_p, p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dv!(dv_m, v_m, .-p.r, ff.x_lgn, ff.dv_temp, p)
#         fun_dx_v1!(dx, x, ff.C, z, p.x_V2, p)
#         fun_dy!(dy, y, ff.C, x, ff.m, ff.dy_temp, p)
#         fun_dm!(dm, ff.m, x, ff.dm_temp, p)
#         fun_dz!(dz, z, y, ff.H_z, ff.s, ff.dz_temp, p)
#         fun_ds!(ds, ff.s, ff.H_z, ff.ds_temp, p)

# # 		@. du[:, :, 1:p.K,:] = ff.dx
# # 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# # 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# # 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# # 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

# #         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# # 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing

# end

# mutable struct LamFunction_gpu2{T <: CuArray{Float32,4,Nothing}} <: Function
# 	x::T
# 	y::T
# 	m::T
# 	z::T
# 	s::T
# 	v_p::T
# 	v_m::T

# 	dx::T
# 	dy::T
# 	dm::T
# 	dz::T
# 	ds::T
# 	dv_p::T
# 	dv_m::T

# 	x_lgn::T
# 	C::T
# 	H_z::T
# 	dy_temp::T
# 	ds_temp::T
# 	H_z_temp::T
# 	V_temp_1::T
# 	V_temp_2::T
# 	Q_temp::T
# 	P_temp::T
# end

# mutable struct LamFunction_cpu2{T <: AbstractArray} <: Function
# 	x::T
# 	y::T
# 	m::T
# 	z::T
# 	s::T
# 	v_p::T
# 	v_m::T

# 	dx::T
# 	dy::T
# 	dm::T
# 	dz::T
# 	ds::T
# 	dv_p::T
# 	dv_m::T

# 	x_lgn::T
# 	C::T
# 	H_z::T
# 	ds_temp::T
# 		H_z_temp::T
# 	V_temp_1::T
# 	V_temp_2::T
# 	Q_temp::T
# 	P_temp::T
# end

# mutable struct LamFunction{T} <: Function
# 	x::AbstractArray
# 	y::AbstractArray
# 	m::AbstractArray
# 	z::AbstractArray
# 	s::AbstractArray
# 	v_p::AbstractArray
# 	v_m::AbstractArray

# 	dx::AbstractArray
# 	dy::AbstractArray
# 	dm::AbstractArray
# 	dz::AbstractArray
# 	ds::AbstractArray
# 	dv_p::AbstractArray
# 	dv_m::AbstractArray

# 	x_lgn::AbstractArray
# 	C::AbstractArray
# 	H_z::AbstractArray
# 	V_temp_1::AbstractArray
# 	V_temp_2::AbstractArray
# 	Q_temp::AbstractArray
# 	P_temp::AbstractArray
# end

# function (ff::LamFunction_cpu2)(du, u, p, t)

#     @inbounds begin
#         @. ff.x = @view u[:, :, 1:p.K,:]
#         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

#         @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

#         fun_x_lgn!(ff.x_lgn, ff.x, p)
#         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

#         fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, p)
#         fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, p)
#         fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
#         fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, p)
#         fun_dm!(ff.dm, ff.m, ff.x, p)
#         fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, p)
#         fun_ds!(ff.ds, ff.s, ff.H_z, ff.ds_temp, p)

# 		@. du[:, :, 1:p.K,:] = ff.dx
# 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

#         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing
# end

# function (ff::LamFunction_gpu2)(du, u, p, t)

#        @inbounds begin
#         @. ff.x = @view u[:, :, 1:p.K,:]
#         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

#         @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

#         fun_x_lgn!(ff.x_lgn, ff.x, p)
#         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, ff.z, ff.H_z_temp, p)

#         fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, p)
#         fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, p)
#         fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
#         fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
#         fun_dm!(ff.dm, ff.m, ff.x, p)
#         fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, p)
#         fun_ds!(ff.ds, ff.s, ff.H_z, ff.ds_temp, p)

# 		@. du[:, :, 1:p.K,:] = ff.dx
# 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

#         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing

# end

# mutable struct LamFunction_04{T} <: Function
# x::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# y::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# m::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# z::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# s::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# v_p::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# v_m::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dx::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dy::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dm::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dz::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# ds::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dv_p::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# dv_m::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	x_lgn::CuArray{Float32,4,Nothing}
# 	C::CuArray{Float32,4,Nothing}
# 	H_z::CuArray{Float32,4,Nothing}
# 	V_temp_1::CuArray{Float32,4,Nothing}
# 	V_temp_2::CuArray{Float32,4,Nothing}
# 	Q_temp::CuArray{Float32,4,Nothing}
# 	P_temp::CuArray{Float32,4,Nothing}
# end

# mutable struct LamFunction_10 <: Function
# 	x
# 	y
# 	m
# 	z
# 	s
# 	v_p
# 	v_m
# 	dx
# 	dy
# 	dm
# 	dz
# 	ds
# 	dv_p
# 	dv_m
# 	x_lgn
# 	C
# 	H_z
# 	V_temp_1
# 	V_temp_2
# 	Q_temp
# 	P_temp
# end

# mutable struct LamFunction_13<: Function
# 	x::CuArray{Float32,4,Nothing}
# 	y::CuArray{Float32,4,Nothing}
# 	m::CuArray{Float32,4,Nothing}
# 	z::CuArray{Float32,4,Nothing}
# 	s::CuArray{Float32,4,Nothing}
# 	v_p::CuArray{Float32,4,Nothing}
# 	v_m::CuArray{Float32,4,Nothing}
# dx::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dy::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dm::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dz::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	ds::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dv_p::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	dv_m::SubArray{Float32,4,CuArray{Float32,4,Nothing},Tuple{Base.Slice{Base.OneTo{Int64}},Base.Slice{Base.OneTo{Int64}},UnitRange{Int64},Base.Slice{Base.OneTo{Int64}}},false}
# 	x_lgn::CuArray{Float32,4,Nothing}
# 	C::CuArray{Float32,4,Nothing}
# 	H_z::CuArray{Float32,4,Nothing}
# 	V_temp_1::CuArray{Float32,4,Nothing}
# 	V_temp_2::CuArray{Float32,4,Nothing}
# 	Q_temp::CuArray{Float32,4,Nothing}
# 	P_temp::CuArray{Float32,4,Nothing}
# end

# function (ff::LamFunction_13)(du, u, p, t)

#     @inbounds begin
# #         ff.x = @view u[:, :, 1:p.K,:]
# #         ff.y = @view u[:, :, p.K+1:2*p.K,:]
# #         ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
# #         ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
# #         ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

# #         ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# # 		ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# 		@. ff.x = @view u[:, :, 1:p.K,:]
#         @. ff.y = @view u[:, :, p.K+1:2*p.K,:]
#         @. ff.m = @view u[:, :, 2*p.K+1:3*p.K,:]
#         @. ff.z = @view u[:, :, 3*p.K+1:4*p.K,:]
#         @. ff.s = @view u[:, :, 4*p.K+1:5*p.K,:]

#         @. ff.v_p = @view u[:, :, 5*p.K+1:5*p.K+1,:]
# 		@. ff.v_m = @view u[:, :, 5*p.K+2:5*p.K+2,:]

# # 		ff.dx = @view du[:, :, 1:p.K,:]
# # 		ff.dy = @view du[:, :, p.K+1:2*p.K,:]
# # 		ff.dm = @view du[:, :, 2*p.K+1:3*p.K,:]
# # 		ff.dz = @view du[:, :, 3*p.K+1:4*p.K,:]
# # 		ff.ds = @view du[:, :, 4*p.K+1:5*p.K,:]

# # 		ff.dv_p = @view du[:, :, 5*p.K+1:5*p.K+1,:]
# # 		ff.dv_m = @view du[:, :, 5*p.K+2:5*p.K+2,:]

# # 		ff.x_lgn = @view ff.x_lgn[:,:,:,:]
# # 		ff.C_ = @view ff.C[:,:,:,:]
# # 		ff.H_z_ = @view ff.H_z[:,:,:,:]
# # 		ff.V_temp_1_ = @view ff.V_temp_1[:,:,:,:]
# # 		ff.V_temp_2_ = @view ff.V_temp_2[:,:,:,:]
# # 		ff.Q_temp_ = @view ff.Q_temp[:,:,:,:]
# # 		ff.P_temp_ = @view ff.P_temp[:,:,:,:]

#         fun_x_lgn!(ff.x_lgn, ff.x, p)
#         fun_v_C!(ff.C, ff.v_p, ff.v_m, ff.V_temp_1, ff.V_temp_2, ff.Q_temp, ff.P_temp, p)
#         fun_H_z!(ff.H_z, ff.z, p)

#         fun_dv!(ff.dv_p, ff.v_p, p.r, ff.x_lgn, p)
#         fun_dv!(ff.dv_m, ff.v_m, .-p.r, ff.x_lgn, p)
#         fun_dx_v1!(ff.dx, ff.x, ff.C, ff.z, p.x_V2, p)
#         fun_dy!(ff.dy, ff.y, ff.C, ff.x, ff.m, ff.dy_temp, p)
#         fun_dm!(ff.dm, ff.m, ff.x, p)
#         fun_dz!(ff.dz, ff.z, ff.y, ff.H_z, ff.s, p)
#         fun_ds!(ff.ds, ff.s, ff.H_z, p)

# # 		@. du[:, :, 1:p.K,:] = ff.dx
# # 		@. du[:, :, p.K+1:2*p.K,:] = ff.dy
# # 		@. du[:, :, 2*p.K+1:3*p.K,:] = ff.dm
# # 		@. du[:, :, 3*p.K+1:4*p.K,:] = ff.dz
# # 		@. du[:, :, 4*p.K+1:5*p.K,:] = ff.ds

# #         @. du[:, :, 5*p.K+1:5*p.K+1,:] = ff.dv_p
# # 		@. du[:, :, 5*p.K+2:5*p.K+2,:] = ff.dv_m

#     end
#     return nothing
# end

function kernels_gpu(img::AbstractArray, p::NamedTuple)
       C_Q_temp = reshape(
        Array{eltype(img)}(undef, p.C_AB_l, p.C_AB_l * p.K),
        p.C_AB_l,
        p.C_AB_l,
        1,
		p.K
    )
C_P_temp = similar(C_Q_temp)
	    H_temp = reshape(
        zeros(eltype(img), p.H_l, p.H_l * p.K * p.K),
        p.H_l,
        p.H_l,
        p.K,
    p.K)
 T_temp = reshape(Array{eltype(img)}(undef, p.K * p.K), 1, 1, p.K, p.K)
 W_temp =
        reshape(Array{eltype(img)}(undef, p.W_l, p.W_l * p.K * p.K), p.W_l, p.W_l, p.K, p.K)
    for k ∈ 1:p.K
        θ = π * (k - 1.0f0) / p.K
        C_Q_temp[:, :, 1,k] = LamKernels.kern_d(p.σ_2, θ)
        C_P_temp[:, :, 1,k] = LamKernels.kern_b(p.σ_2, θ)
        H_temp[:, :, k,k] = p.H_fact .* LamKernels.gaussian_rot(p.H_σ_x, p.H_σ_y, θ, p.H_l)
# 		todo make T kernel more general for higher K
        T_temp[1, 1, k,1] = p.T_fact[k]
        T_temp[1, 1, 2,2] = p.T_fact[1]
        T_temp[1, 1, 1,2] = p.T_fact[2]
        #todo: generalise T and W for higher K
        #         T_temp[:,:,k] = KernelFactors.gaussian(p.T_σ, p.K)
        #         for l ∈ 1:p.K
        #             W_temp[:,:,l,k] =
        #         end
    end

    W_temp[:, :, 1, 1] =
        5f0 .* LamKernels.gaussian_rot(3f0, 0.8f0, 0f0, p.W_l) .+
        LamKernels.gaussian_rot(0.4f0, 1f0, 0f0, p.W_l)

	W_temp[:, :, 2, 2] =
        5f0 .* LamKernels.gaussian_rot(3f0, 0.8f0,  π / 2f0, p.W_l) .+
        LamKernels.gaussian_rot(0.4f0, 1f0, π / 2f0, p.W_l)

	W_temp[:, :, 1, 2] = relu.(
        0.2f0 .- LamKernels.gaussian_rot(2f0, 0.6f0, 0f0, p.W_l) .-
        LamKernels.gaussian_rot(0.3f0, 1.2f0, 0f0, p.W_l))

	W_temp[:, :, 2, 1] = relu.(
        0.2f0 .- LamKernels.gaussian_rot(2f0, 0.6f0, π / 2f0, p.W_l) .-
        LamKernels.gaussian_rot(0.3f0, 1.2f0, π / 2f0, p.W_l))

temp_out = (
        k_gauss_1 = cu(reshape2d_4d(Kernel.gaussian(p.σ_1))),
        k_gauss_2 = cu(reshape2d_4d(Kernel.gaussian(p.σ_2))),
        k_C_d = cu(C_Q_temp),
        k_C_b = cu(C_P_temp),

# 		todo use mean of x_lgn?
		k_x_lgn = cu(reshape(ones(Float32,1,p.K),1,1,p.K,1)),
# 		k_x_lgn = cu(reshape(ones(Float32,1,p.K),1,1,p.K,1) ./ p.K),
        k_W_p = cu(W_temp),
        k_W_m = cu(W_temp),
        k_H = cu(H_temp),
        k_T_p = cu(T_temp),
        k_T_m = cu(p.T_p_m .* T_temp),
        k_T_p_v2 = cu(p.T_v2_fact .* T_temp),
        k_T_m_v2 = cu(p.T_v2_fact .* p.T_p_m .* T_temp),
        dim_i = size(img)[1],
        dim_j = size(img)[2],
        x_V2 = cu(reshape(zeros(Float32, size(img)[1], size(img)[2] * p.K), size(img)[1], size(img)[2],p.K,1)),
ν_pw_n= p.ν^p.n, )

merge(p, temp_out)
end

function kernels_cpu(img::AbstractArray, p::NamedTuple)
       C_Q_temp = reshape(
        Array{eltype(img)}(undef, p.C_AB_l, p.C_AB_l * p.K),
        p.C_AB_l,
        p.C_AB_l,
        1,
		p.K
    )
C_P_temp = similar(C_Q_temp)
	    H_temp = reshape(
        zeros(eltype(img), p.H_l, p.H_l * p.K * p.K),
        p.H_l,
        p.H_l,
        p.K,
    p.K)
 T_temp = reshape(Array{eltype(img)}(undef, p.K * p.K), 1, 1, p.K, p.K)
 W_temp =
        reshape(Array{eltype(img)}(undef, p.W_l, p.W_l * p.K * p.K), p.W_l, p.W_l, p.K, p.K)
    for k ∈ 1:p.K
        θ = π * (k - 1.0f0) / p.K
        C_Q_temp[:, :, 1,k] = LamKernels.kern_d(p.σ_2, θ)
        C_P_temp[:, :, 1,k] = LamKernels.kern_b(p.σ_2, θ)
        H_temp[:, :, k,k] = p.H_fact .* LamKernels.gaussian_rot(p.H_σ_x, p.H_σ_y, θ, p.H_l)
# 		todo make T kernel more general for higher K
        T_temp[1, 1, k,1] = p.T_fact[k]
        T_temp[1, 1, 2,2] = p.T_fact[1]
        T_temp[1, 1, 1,2] = p.T_fact[2]
        #todo: generalise T and W for higher K
        #         T_temp[:,:,k] = KernelFactors.gaussian(p.T_σ, p.K)
        #         for l ∈ 1:p.K
        #             W_temp[:,:,l,k] =
        #         end
    end

    W_temp[:, :, 1, 1] =
        5f0 .* LamKernels.gaussian_rot(3f0, 0.8f0, 0f0, p.W_l) .+
        LamKernels.gaussian_rot(0.4f0, 1f0, 0f0, p.W_l)

	W_temp[:, :, 2, 2] =
        5f0 .* LamKernels.gaussian_rot(3f0, 0.8f0,  π / 2f0, p.W_l) .+
        LamKernels.gaussian_rot(0.4f0, 1f0, π / 2f0, p.W_l)

	W_temp[:, :, 1, 2] = relu.(
        0.2f0 .- LamKernels.gaussian_rot(2f0, 0.6f0, 0f0, p.W_l) .-
        LamKernels.gaussian_rot(0.3f0, 1.2f0, 0f0, p.W_l))

	W_temp[:, :, 2, 1] = relu.(
        0.2f0 .- LamKernels.gaussian_rot(2f0, 0.6f0, π / 2f0, p.W_l) .-
        LamKernels.gaussian_rot(0.3f0, 1.2f0, π / 2f0, p.W_l))

	temp_out = (
        k_gauss_1 = reshape2d_4d(Kernel.gaussian(p.σ_1)),
        k_gauss_2 = reshape2d_4d(Kernel.gaussian(p.σ_2)),
        k_C_d = C_Q_temp,
        k_C_b = C_P_temp,

# 		todo use mean of x_lgn?
		k_x_lgn = reshape(ones(Float32,1,p.K),1,1,p.K,1),
# 		k_x_lgn = reshape(ones(Float32,1,p.K),1,1,p.K,1)./p.K,
        k_W_p = W_temp,
        k_W_m = W_temp,
        k_H = H_temp,
        k_T_p = T_temp,
        k_T_m = p.T_p_m .* T_temp,
        k_T_p_v2 = p.T_v2_fact .* T_temp,
        k_T_m_v2 = p.T_v2_fact .* p.T_p_m .* T_temp,
        dim_i = size(img)[1],
        dim_j = size(img)[2],
        x_V2 = reshape(zeros(Float32, size(img)[1], size(img)[2] * p.K), size(img)[1], size(img)[2],p.K,1),
ν_pw_n= p.ν^p.n, )
merge(p, temp_out)
end

function reshape2d_4d(img::AbstractArray)
    reshape(img, size(img)[1], size(img)[2], 1, 1)
end

# function reshape_ijk_ij1k(img::AbstractArray, p::NamedTuple)
#     reshape(img, size(img)[1], size(img)[2], 1, p.K)
# end

function conv!(out::AbstractArray, img::AbstractArray, kern::AbstractArray, p::NamedTuple)
# 	out_ = @view out[:,:,:,:]
# 	img_ = @view img[:,:,:,:]
# 	kern_ = @view kern[:,:,:,:]
# 	    @inbounds out .= NNlib.conv(img[:,:,:,:], kern, pad=(size(kern)[1]>>1, size(kern)[1]>>1, size(kern)[2]>>1, size(kern)[2]>>1), flipped=true)

#     @inbounds out .= NNlib.conv(img, kern, pad=(size(kern)[1]>>1, size(kern)[1]>>1, size(kern)[2]>>1, size(kern)[2]>>1), flipped=true)
# 	    @inbounds out .= NNlib.conv(img_, kern_, pad=(size(kern_)[1]>>1, size(kern_)[1]>>1, size(kern_)[2]>>1, size(kern_)[2]>>1), flipped=true)
    @inbounds NNlib.conv!(out, img, kern, NNlib.DenseConvDims(img, kern, padding = ( size(kern)[1]>>1 , size(kern)[1]>>1 , size(kern)[2]>>1 , size(kern)[2]>>1 ), flipkernel = true))

# 	@. out_ = out
    return nothing
end

# function conv!(out::AbstractArray, img::AbstractArray, kern::AbstractArray, p::NamedTuple)
# # 	out_ = @view out[:,:,:,:]
#     @inbounds out = NNlib.conv(img, kern, pad=(size(kern)[1]>>1, size(kern)[1]>>1, size(kern)[2]>>1, size(kern)[2]>>1), flipped=true)
# # 	@. out_ = out
#     return nothing
# end

function add_I_u_p(I::AbstractArray, p::NamedTuple)
# 	todo fix
# 	I_4d = cu(reshape2d_4d(I))
	I_4d = reshape2d_4d(I)
	r = similar(I_4d)
	I_u!(r, I_4d, p)
    temp_out = (I = I_4d, r = r)
    return merge(p, temp_out)
end

# retina

function I_u!(r::AbstractArray, I::AbstractArray, p::NamedTuple)
	conv!(r, I, p.k_gauss_1, p)
    @. r = I - r
end

# lgn feedback

# function fun_x_lgn!(x_lgn::AbstractArray, x::AbstractArray, p::NamedTuple)
# # 	out_ = @view x_lgn[:,:,:,:]
# 	    @inbounds x_lgn .= NNlib.conv(x[:,:,:,:], p.k_x_lgn, pad=0,flipped=true)
#     return nothing
# end

function fun_x_lgn!(x_lgn::AbstractArray, x::AbstractArray, p::NamedTuple)
# 	out_ = @view x_lgn[:,:,:,:]
# 	x_ = @view x[:,:,:,:]
# 	kern_ = @view p.k_x_lgn[:,:,:,:]
# 	    @inbounds x_lgn .= NNlib.conv(x_, p.k_x_lgn, pad=0,flipped=true)

		    @inbounds NNlib.conv!(x_lgn, x, p.k_x_lgn, NNlib.DenseConvDims(x, p.k_x_lgn, padding=0, flipkernel=true))

    return nothing
end
# function fun_F(value::Real, p::NamedTuple)
#     max.(value - p.Γ, 0)
# end
#
#
# # williomson uses differnt F, relu with threshold
# function fun_F_willimson(value::Real, p::NamedTuple)
#     value < p.Γ ? zero(value) : value
# end

function fun_f!(arr::AbstractArray, p::NamedTuple)
   @inbounds @. arr = (p.μ * arr^p.n) / (p.ν_pw_n + arr^p.n)
    return nothing
end

# LGN
function fun_dv!(
    dv::AbstractArray,
    v::AbstractArray,
    u::AbstractArray,
    x_lgn::AbstractArray,
	dv_temp::AbstractArray,
    p::NamedTuple,
)
	@inbounds begin
	conv!(dv_temp, x_lgn, p.k_gauss_1, p)
    @. dv =
        p.δ_v * (
            -v + ((1f0 - v) * max(u, 0f0) * (1f0 + p.C_1 * x_lgn)) -
            ((1f0 + v) * p.C_2 * dv_temp)
        )
	end
    return nothing
end

# lgn to l6 and l4
# function fun_v_C!(
#     v_C::AbstractArray,
#     v_p::AbstractArray,
#     v_m::AbstractArray,
# 	V_temp_1::AbstractArray,
# 	V_temp_2::AbstractArray,
# 	Q_temp::AbstractArray,
# 	P_temp::AbstractArray,
#     p::NamedTuple,
# )
# #     V = similar(v_p)
# #     temp = similar(v_p)

# # 	A = similar(v_C)
# # 	B = similar(v_C)
#     #     allocate B to v_C

# # 	mx_v_p = similar(v_p)

# # 	@. mx_v_p = max(v_p,0f0)

#     @inbounds @. V_temp_2 = exp(-1.0f0 / 8.0f0) * (max(v_p[:,:,:,:], 0f0) - max(v_m[:,:,:,:], 0f0))
#     conv!(V_temp_1, V_temp_2, p.k_gauss_2, p)

# 	conv!(Q_temp, V_temp_1, p.k_C_d, p)
# 	conv!(P_temp, V_temp_1, p.k_C_b, p)
# # 	@. P_temp = abs(v_C)
# 	@inbounds @. P_temp = abs(v_C[:,:,:,:])

#     @inbounds @. v_C = p.γ * (max(Q_temp[:,:,:,:] - P_temp[:,:,:,:], 0f0) + max(-Q_temp[:,:,:,:] - P_temp[:,:,:,:], 0f0))

#     return nothing
# end

function fun_v_C!(
    v_C::AbstractArray,
    v_p::AbstractArray,
    v_m::AbstractArray,
	V_temp_1::AbstractArray,
	V_temp_2::AbstractArray,
	Q_temp::AbstractArray,
	P_temp::AbstractArray,
    p::NamedTuple,
)
#     V = similar(v_p)
#     temp = similar(v_p)

# 	A = similar(v_C)
# 	B = similar(v_C)
    #     allocate B to v_C

# 	mx_v_p = similar(v_p)

# 	@. mx_v_p = max(v_p,0f0)

    @inbounds begin
		@. V_temp_2 = exp(-1.0f0 / 8.0f0) * (max(v_p, 0f0) - max(v_m, 0f0))
    conv!(V_temp_1, V_temp_2, p.k_gauss_2, p)

	conv!(Q_temp, V_temp_1, p.k_C_d, p)
	conv!(P_temp, V_temp_1, p.k_C_b, p)
	@. P_temp = abs(P_temp)

    @. v_C = p.γ * (max(Q_temp - P_temp, 0f0) + max(-Q_temp - P_temp, 0f0))
	end

    return nothing
end

# L6
function fun_dx_v1!(
    dx::AbstractArray,
    x::AbstractArray,
    C::AbstractArray,
    z::AbstractArray,
    x_v2::AbstractArray,
    p::NamedTuple,
)
	@inbounds begin
    @. dx =
        p.δ_c * (
            -x + (
                (1.0f0 - x) * (
                    (p.α * C) + (p.ϕ * max(z - p.Γ, 0f0)) .+ (p.v_21 * x_v2) +
                    p.att
                )
            )
        )
	end
    return nothing
end

#     L4 excit
function fun_dy!(
    dy::AbstractArray,
    y::AbstractArray,
    C::AbstractArray,
    x::AbstractArray,
    m::AbstractArray,
	dy_temp::AbstractArray,
    p::NamedTuple,
)
	@inbounds begin
    conv!(dy_temp, m, p.k_W_p, p)
    @. dy_temp = m * dy_temp
    fun_f!(dy_temp, p)
    @. dy = p.δ_c * (-y + ((1f0 - y) * (C + (p.η_p * x))) - ((1f0 + y) * dy_temp))
	end
    return nothing
end

#     l4 inhib
function fun_dm!(
    dm::AbstractArray,
    m::AbstractArray,
    x::AbstractArray,
	dm_temp::AbstractArray,
    p::NamedTuple,
)
	@inbounds begin
    conv!(dm_temp, m, p.k_W_m, p)
    fun_f!(dm_temp, p)
    @. dm = p.δ_m * (-m + (p.η_m * x) - (m * dm_temp))
	end
    return nothing
end

#     L2/3 excit
function fun_dz!(
    dz::AbstractArray,
    z::AbstractArray,
    y::AbstractArray,
    H_z::AbstractArray,
    s::AbstractArray,
	dz_temp::AbstractArray,
    p::NamedTuple,
)
	@inbounds begin
    conv!(dz_temp, s, p.k_T_p, p)
    @. dz =
        p.δ_z * (
            -z + ((1f0 - z) * ((p.λ * max(y, 0f0)) + H_z + (p.a_23_ex * p.att))) -
            ((z + p.ψ) * dz_temp)
        )
	end
    return nothing
end

#     L2/3 inhib
function fun_ds!(
    ds::AbstractArray,
    s::AbstractArray,
    H_z::AbstractArray,
		ds_temp::AbstractArray,
    p::NamedTuple,
)
	@inbounds begin
    conv!(ds_temp, s, p.k_T_m, p)
    @. ds = p.δ_s * (-s + H_z + (p.a_23_in * p.att) - (s * ds_temp))
		end
	return nothing
end

function fun_H_z!(
		H_z::AbstractArray,
		z::AbstractArray,
		H_z_temp::AbstractArray,
		p::NamedTuple)
	@inbounds begin
	@. H_z_temp = max(z - p.Γ, 0f0)
	conv!(H_z, H_z_temp, p.k_H, p)
	end
    return nothing
end

#     V2 L6
function fun_dx_v2!(
    dx::AbstractArray,
    x_v2::AbstractArray,
    z_v2::AbstractArray,
    z::AbstractArray,
    p::NamedTuple,
)
	@inbounds begin
    @. dx =
        p.δ_c * (
            -x_v2 + (
                (1f0 - x_v2) * (
                    (p.v12_6 * max(z - p.Γ, 0f0)) +
                    (p.ϕ * max(z_v2 - p.Γ, 0f0)) +
                    p.att
                )
            )
        )
	end
    return nothing
end

# V2 L4 excit
function fun_dy_v2!(
    dy::AbstractArray,
    y_v2::AbstractArray,
    z::AbstractArray,
    x_v2::AbstractArray,
    m_v2::AbstractArray,
		dy_temp::AbstractArray,
    p::NamedTuple,
)
	@inbounds begin
    conv!(dy_temp, m_v2, p.k_W_p, p)
    fun_f!(dy_temp, p)
    @. dy =
        p.δ_c * (
            -y_v2 +
            ((1.0f0 - y_v2) * ((p.v12_4 * max(z - p.Γ, 0f0)) + (p.η_p * x_v2))) -
            ((1.0f0 + y_v2) * dy_temp)
        )
	end
    return nothing
end

# # # for equilabrium

# # l6 equilabrium
function fun_x_equ!(x::AbstractArray, C::AbstractArray, z::AbstractArray, x_v2::AbstractArray, x_temp::AbstractArray, p::NamedTuple)
	@inbounds begin
		@. x_temp = (p.α * C) + (p.ϕ * max(z - p.Γ, 0f0)) + (p.v_21 * x_v2) + p.att
		@. x = x_temp /(1+x_temp)
	end
    return nothing
end

# # l4 excit equilabrium
function fun_y_equ!(y::AbstractArray, C::AbstractArray, x::AbstractArray,  m::AbstractArray, dy_temp::AbstractArray, p::NamedTuple)
	@inbounds begin
		conv!(dy_temp, m, p.k_W_p, p)
		@. dy_temp = m * dy_temp
		fun_f!(dy_temp, p)
		@. y = (C + (p.η_p * x) - dy_temp)/(1 + C + (p.η_p * x) + dy_temp)
	end
    return nothing
end

# # lgn
# function fun_v_equ(u::AbstractArray, x::AbstractArray, p::NamedTuple)
#     return fun_equ.(relu.(u) .* (p.lgn_equ_u .+ (p.lgn_equ_A .* fun_A.(x))) .- (p.lgn_equ_B .* fun_B.(x)))
# end
#
#
# # lgn, no L6 feedback, light
# function fun_v_equ_noFb(u::AbstractArray, x::AbstractArray, p::NamedTuple)
#     return fun_equ.(relu.(u) .* (p.lgn_para_u))
# end
#
#
# # l6
# function fun_x_equ(C::AbstractArray, z::AbstractArray,  p::NamedTuple)
#     return fun_equ.((p.α .* C) .+ (p.ϕ .* fun_F.(z,p.Γ)) + (p.V_21 .* x_v2) .+ p.att)
# end
#
#
# # l6, no V2 feedback, light
# function fun_x_equ_noV2(C::AbstractArray, z::AbstractArray, p::NamedTuple)
#     return fun_equ.((p.α .* C) .+ (p.ϕ .* fun_F.(z,p.Γ)))
# end
#
#
# # l4 excit
# function fun_y_equ(C::AbstractArray, x::AbstractArray,  m::AbstractArray, p::NamedTuple)
#     return fun_equ.(C .+ (η_p .* x - fun_f.(imfilter(m, p.k_W_p), p.μ, p.ν, p.n)))
# end
#
# # l4 inhib, needs initial condition of itself
# function fun_m_equ(C::AbstractArray, x::AbstractArray, m_init::AbstractArray, p::NamedTuple)
#     return (p.η_m .* x ./ (1 .+  fun_f.(imfilter(m_init, p.k_W_m), p.μ, p.ν, p.n)))
# end
#
#
# # l4 inhib - no feedback kernel
# function fun_m_equ_noFb(C::AbstractArray, x::AbstractArray, p::NamedTuple)
#     return (p.η_m .* x)
# end
#
#
# #  l2/3 excit, needs initial condition of itself
# function fun_z_equ(y::AbstractArray, s::AbstractArray,  z_init::AbstractArray,  p::NamedTuple)
#     return (p.λ .* relu.(y)) .+ imfilter(fun_F.(z_init, p.Γ), k_H) .+ (p.a_23_ex .* p.att) .- (p.ϕ .* imfilter(s, k_T_p)) ./
#     (1 .+ (p.λ .* relu.(y)) .+ imfilter(fun_F.(z_init, p.Γ), k_H) .+ (p.a_23_ex .* p.att) .+ imfilter(s, k_T_p))
#     return
# end
#
#
# #  l2/3 excit - no feedback kernel
# function fun_z_equ_noFb(y::AbstractArray, s::AbstractArray, p::NamedTuple)
#     return (p.λ .* relu.(y)) .- (p.ϕ .* imfilter(s, k_T_p)) ./
#     (1 .+ (p.λ .* relu.(y)) .+ imfilter(s, k_T_p))
#     return
# end
#
#
# #  l2/3 inhib, needs initial condition of itself
# function fun_s_equ(z::AbstractArray, s::AbstractArray,  s_init::AbstractArray, p::NamedTuple)
#     return ((imfilter(fun_F.(z, p.Γ), k_H) + p.a_23_in .* p.att ) ./ (1 .+ imfilter(s_init, k_T_m)))           #????????? is T right?
# end
#
#
# #  l2/3 inhib - no feedback kernel
# function fun_s_equ_noFb(z::AbstractArray, s::AbstractArray, p::NamedTuple)
#     return (imfilter(fun_F.(z, p.Γ), k_H))
# end
#
#
# # V2 L6
# function fun_xV2_equ(x_v2::AbstractArray, z::AbstractArray, z_v2::AbstractArray, p::NamedTuple)
#     return fun_equ.((p.V_12_6 .* fun_F.(z, p.Γ) .+ (p.ϕ .* fun_F.(z_v2, p.Γ))))
# end
#
#
# #  V2 L4 excit
# function fun_yV2_equ(z::AbstractArray, x::AbstractArray,  m::AbstractArray, p::NamedTuple)
#     return (p.V_12_4 .* fun_F.(z,p.Γ) .+ (p.η_p .* x) - fun_f.(imfilter(m, k_W_p), p.μ, p.ν, p.n) ./
#    (1 .+ p.V_12_4 .* fun_F.(z,p.Γ) .+ (p.η_p .* x) .+ fun_f.(imfilter(m, k_W_p), p.μ, p.ν, p.n)))
# end

end

# function varables(I::AbstractArray, p::NamedTuple)
#    v_p = zeros(typeof(I[1,1]), size(I)[1], size(I)[2])
#    v_m = copy(v_p)
#    x_lgn = copy(v_p)
#    x = reshape(zeros(typeof(I[1,1]), size(I)[1], size(I)[2] * p.K), size(I)[1], size(I)[2], p.K)
#    y = copy(x)
#    m = copy(x)
#    z = copy(x)
#    s = copy(x)
#    C = copy(x)
#    H_z = copy(x)
#    x_V2 = copy(x)
#    y_V2 = copy(x)
#    m_V2 = copy(x)
#    z_V2 = copy(x)
#    s_V2 = copy(x)
#    H_z_V2 = copy(x)
#    return [v_p, v_m, x_lgn, x, y, m, z, s, C, H_z, x_V2, y_V2, m_V2, z_V2, s_V2, H_z_V2]
# end
#
#
# function variables_sep(I::AbstractArray, p::NamedTuple)
#    v_p = zeros(typeof(I[1,1]), size(I)[1], size(I)[2])
#    v_m = copy(v_p)
#    x_lgn = copy(v_p)
#    x = reshape(zeros(typeof(I[1,1]), size(I)[1], size(I)[2] * p.K), size(I)[1], size(I)[2], p.K)
#    y = copy(x)
#    m = copy(x)
#    z = copy(x)
#    s = copy(x)
#    C = copy(x)
#    H_z = copy(x)
#    x_V2 = copy(x)
#    y_V2 = copy(x)
#    m_V2 = copy(x)
#    z_V2 = copy(x)
#    s_V2 = copy(x)
#    H_z_V2 = copy(x)
#    return v_p, v_m, x_lgn, x, y, m, z, s, C, H_z, x_V2, y_V2, m_V2, z_V2, s_V2, H_z_V2
# end
#
# function variables_dict(I::AbstractArray, p::NamedTuple)
#    v_p = zeros(typeof(I[1,1]), size(I)[1], size(I)[2])
#    v_m = copy(v_p)
#    x_lgn = copy(v_p)
#    x = reshape(zeros(typeof(I[1,1]), size(I)[1], size(I)[2] * p.K), size(I)[1], size(I)[2], p.K)
#    y = copy(x)
#    m = copy(x)
#    z = copy(x)
#    s = copy(x)
#    C = copy(x)
#    H_z = copy(x)
#    x_V2 = copy(x)
#    y_V2 = copy(x)
#    m_V2 = copy(x)
#    z_V2 = copy(x)
#    s_V2 = copy(x)
#    H_z_V2 = copy(x)
#    return v_p, v_m, x_lgn, x, y, m, z, s, C, H_z, x_V2, y_V2, m_V2, z_V2, s_V2, H_z_V2
# end

# function varables(I::AbstractArray, p::NamedTuple)
#    ij = zeros(typeof(I[1,1]), size(I)[1], size(I)[2])
#    ijk = reshape(zeros(typeof(I[1,1]), size(I)[1], size(I)[2] * p.K), size(I)[1], size(I)[2], p.K)
#    temp_out = [v_p = copy(ij),
#    v_m = copy(ij),
#    x_lgn = copy(ij),
#    x = copy(ijk),
#    y = copy(ijk),
#    m = copy(ijk),
#    z = copy(ijk),
#    s = copy(ijk),
#    C = copy(ijk),
#    H_z = copy(ijk),
#    x_V2 = copy(ijk),
#    y_V2 = copy(ijk),
#    m_V2 = copy(ijk),
#    z_V2 = copy(ijk),
#    s_V2 = copy(ijk),
#    H_z_V2 = copy(ijk)]
#    return temp_out
# end


text:
using DiscreteElasticRods
using BenchmarkTools
using LinearAlgebra
using ForwardDiff
using SparseDiffTools
using IterativeSolvers

ns = 50

# Typing for dual caches
DualTag = SparseDiffTools.DeivVecTag
DualType = ForwardDiff.Dual{DualTag, Float64, 1}

# For now, all dual_caches are tuples
function allocate_dual_cache(F::Function, DType::Type{ForwardDiff.Dual{T,V,N}}, Ns) where{T,V,N}
    (DER.allocate_cache(F, V, Ns), DER.allocate_cache(F, DType, Ns))
end

function allocate_dual_vector(DType::Type{ForwardDiff.Dual{T,V,N}}, Ne) where{T,V,N}
    (zeros(V, Ne), zeros(DType, Ne))
end

function allocate_dual_delta(DType::Type{ForwardDiff.Dual{T,V,N}}, Ns) where{T,V,N}
    (DER.allocate_delta(V, Ns), DER.allocate_delta(DType, Ns))
end

resolve_type(x, ::Type) = x[1]
function resolve_type(x, ::Type{ForwardDiff.Dual{T,V,N}}) where{T,V,N}
    x[2]
end

# Define a Rod
p1 = [0, 0, 0]
p2 = [1, 0, 0]
R = 1.0
d = [0, 0, 1.0]

ref_rod = DER.arc_rod(p1, p2, d, R, ns)
ref_strains = DER.full_kinematics(ref_rod)
ref_strains.κ .= 0.0
ref_strains.τ .= 0.0
rod_props = DER.elliptical_stiffness(1e6, 0.48, 0.01, 0.01)
DER.plot(ref_rod)

# Allocate coordinate layer
s = zeros(Float64, 4*ns)
∇s = zeros(Float64, 4*ns)

# Allocate expanded coordinate layer
q = allocate_dual_vector(DualType, 4*ns+7)
∇q = allocate_dual_vector(DualType, 4*ns+7)

# Convert between the two
function s_to_q!(q, s)
    q[8:end] .= s
end

function ∇q_to_∇s!(∇s, ∇q)
    ∇s .= @view ∇q[8:end]
end

# Non-allocating gradient for Float64
gradients = allocate_dual_delta(DualType, ns)
coordinates = allocate_dual_delta(DualType, ns)
caches = allocate_dual_cache(DER.fused_update_gradient!, DualType, ns)

function dEdq!(∇q, q)
    T = eltype(q)
    coord = resolve_type(coordinates,T)
    grads = resolve_type(gradients,T)
    cache = resolve_type(caches,T)

    DER.banded_storage!(coord, q)
    DER.fused_update_gradient!(grads, ref_rod, coord, rod_props, ref_strains, cache)
    DER.banded_storage!(∇q, grads)
end

function dEds!(∇s, s)
    q_res = resolve_type(q, eltype(s))
    ∇q_res = resolve_type(∇q, eltype(s))

    s_to_q!(q_res, s)
    dEdq!(∇q_res, q_res)
    ∇q_to_∇s!(∇s, ∇q_res)
end

# Test that everything works
s = rand(4*ns)
∇s = zeros(4*ns)
dEds!(∇s, s)

t = DualType.(rand(4*ns))
∇t = DualType.(zeros(4*ns))
dEds!(∇t, t)

@btime dEds!(∇s, s)
@btime dEds!(∇t, t)

# Test for correctness?
using FiniteDiff
newrods = (DER.allocate_rod(Float64, ns), DER.allocate_rod(DualType, ns))
newstrains = (DER.allocate_strain(Float64, ns), DER.allocate_strain(DualType, ns))
function E(s)
    T = eltype(s)
    q_res = resolve_type(q, T)
    coord = resolve_type(coordinates, T)
    cache = resolve_type(caches, T)
    newrod = resolve_type(newrods, T)
    newstrain = resolve_type(newstrains, T)
    s_to_q!(q_res, s)

    DER.banded_storage!(coord, q_res)
    DER.rod_update!(newrod, ref_rod, coord, cache[2])
    DER.full_kinematics!(newstrain, newrod, (cache[2], cache[3]))
    DER.elastic_energy(ref_strains, newstrain, rod_props)
end

∇s0 = similar(s)
dEds!(∇s0, s)

∇s1 = similar(s)
v = zeros(length(s))
for i = 1:length(s)
    v[i] = 1.0
    t .= ForwardDiff.Dual{DualTag}.(s,v)
    Et = E(t)
    ∇s1[i] = ForwardDiff.partials(Et,1)
    v[i] = 0.0
end

using FiniteDiff
∇s2 = FiniteDiff.finite_difference_gradient(E,s)

norm(∇s0 - ∇s1)/norm(∇s1)
norm(∇s1 - ∇s2)/norm(∇s2)
# Passed, move on

# Hessian Vector Products w/ SparseDiffTools
H = SparseDiffTools.JacVec(dEds!, s)
Hdense = zeros(length(s), length(s))
v = zeros(length(s))
for i = 1:length(s)
    v[i] = 1.0
    mul!(view(Hdense, :, i), H, v)
    v[i] = 0.0
end

out1, out2 = minres!(update, H, ∇s; log = true)


text:
"""
    A structure defining a thin film or layer of a Material.
"""
struct Film
    material::Material
    thickness::Float64

    Film() = new(pure(n"C"), 0.0)

    function Film(mat::Material, thickness::AbstractFloat)
        @assert haskey(mat.properties, :Density) "Missing the :Density property when constructing a Film."
        return new(mat, thickness)
    end
end

Base.isequal(f1::Film, f2::Film) =
    isequal(f1.material, f2.material) && isequal(f1.thickness, f2.thickness)

Base.isapprox(f1::Film, f2::Film) =
    isequal(f1.material, f2.material) && isapprox(f1.thickness, f2.thickness, rtol = 1.0e5)

Base.show(io::IO, flm::Film) =
    flm.thickness < 1.0e-4 ?
    print(io, 1.0e7 * flm.thickness, " nm of ", name(flm.material)) : #
    print(io, 1.0e4 * flm.thickness, " μm of ", name(flm.material))

"""
    massthickness(flm::Film)

The mass-thickness of the film in g/cm².
"""
massthickness(flm::Film) = flm.thickness * flm.material[:Density]

"""
    transmission(flm::Film, xrayE::AbstractFloat, θ::AbstractFloat, alg::Type{<:NeXLAlgorithm}=FFASTDB) =

Compute the transmission fraction of an X-ray at the specified angle through a Film.
"""
transmission(
    flm::Film,
    xrayE::AbstractFloat,
    θ::AbstractFloat,
    alg::Type{<:NeXLAlgorithm} = FFASTDB,
) =
    flm.thickness > 0.0 ?
    exp(-mac(flm.material, xrayE, alg) * csc(θ) * massthickness(flm)) : 1.0
transmission(flm::Film, xrayE::AbstractFloat, alg::Type{<:NeXLAlgorithm} = FFASTDB) =
    flm.thickness > 0.0 ? exp(-mac(flm.material, xrayE, alg) * massthickness(flm)) : 1.0

"""
    transmission(flm::Film, cxr::CharXRay, θ::AbstractFloat) =

Compute the transmission fraction of an X-ray at the specified angle through a Film.
"""
transmission(flm::Film, cxr::CharXRay, θ::AbstractFloat) = transmission(flm, energy(cxr), θ)

material(film::Film) = film.material
thickness(film::Film) = film.thickness


text:
struct PDBijector <: Bijector{2} end

# This function has custom adjoints defined for Tracker, Zygote and ReverseDiff.
# I couldn't find a mutation-free implementation that maintains TrackedArrays in Tracker
# and ReverseDiff, hence the need for custom adjoints.
function replace_diag(f, X)
    g(i, j) = ifelse(i == j, f(X[i, i]), X[i, j])
    return g.(1:size(X, 1), (1:size(X, 2))')
end
(b::PDBijector)(X::AbstractMatrix{<:Real}) = pd_link(X)
function pd_link(X)
    Y = lower(parent(cholesky(X; check = true).L))
    return replace_diag(log, Y)
end
(b::PDBijector)(X::AbstractArray{<:AbstractMatrix{<:Real}}) = map(b, X)
lower(A::AbstractMatrix) = convert(typeof(A), LowerTriangular(A))

function (ib::Inverse{<:PDBijector})(Y::AbstractMatrix{<:Real})
    X = replace_diag(exp, Y)
    return getpd(X)
end
(ib::Inverse{<:PDBijector})(X::AbstractArray{<:AbstractMatrix{<:Real}}) = map(ib, X)
getpd(X) = LowerTriangular(X) * LowerTriangular(X)'

function logabsdetjac(b::PDBijector, X::AbstractMatrix{<:Real})
    T = eltype(X)
    Xcf = cholesky(X, check = false)
    if !issuccess(Xcf)
        Xcf = cholesky(X + max(eps(T), eps(T) * norm(X)) * I)
    end
    return logabsdetjac(b, Xcf)
end

function logabsdetjac(b::PDBijector, Xcf::Cholesky)
    U = Xcf.U
    T = eltype(U)
    d = size(U, 1)
    return - sum((d .- (1:d) .+ 2) .* log.(diag(U))) - d * log(T(2))
end

logabsdetjac(b::PDBijector, X::AbstractArray{<:AbstractMatrix{<:Real}}) = mapvcat(X) do x
    logabsdetjac(b, x)
end


text:
using Test
#@startexample Generalised BPS (Truncated Gaussian)
#=

This example is identical to the Global Bouncy Particle Sampler example with
a truncated gaussian except that it uses the *Generalised* kernel. In
that case, there does not need to be a refreshment step.
See also the corresponding paper: [*Generalized Bouncy Particle Sampler*](https://arxiv.org/pdf/1706.04781.pdf) by Changye Wu and Christian
Robert.

=#
using PDSampler, Random
p     = 2
ns, a = diagm(0=>ones(p)), zeros(p)
geom  = Polygonal(ns, a)

nextbdG(x, v) = nextboundary(geom, x, v)

Random.seed!(12)
P1  = randn(p,p)
P1 *= P1'
P1 += norm(P1)/100*diagm(0=>ones(p))
C1  = inv(P1); C1 += C1'; C1/=2;
L1  = cholesky(C1)
mu  = zeros(p) .+ 1.
mvg = MvGaussianCanon(mu, P1)

gradllG(x) = gradloglik(mvg, x)

nextevG(x, v) = nextevent_bps(mvg, x, v)
#=
Note the specification of `algname="GBPS"`. Note also the refreshment rate
set to 0.0 (no refreshment).
=#
T    = 1000.0   # length of path generated
lref = 0.0      # rate of refreshment
x0   = mu+L1.L*randn(p) # sensible starting point
v0   = randn(p) # starting velocity
v0  /= norm(v0) # put it on the sphere (not necessary)
# Define a simulation
sim = Simulation( x0=x0, v0=v0, T=T, nextevent=nextevG, gradloglik=gradllG,
                  nextboundary=nextbdG, lambdaref=lref, maxgradeval = 10000,algname="GBPS")
#=
The rest is as before:
=#
(path, details) = simulate(sim)

sN = 1000
s  = broadcast(+, mu, L1.L*randn(p,sN))
mt = zeros(2)
np = 0
# Sum for all samples in the positive orthan
ss = [s; ones(sN)']
mt = sum(ss[:,i] for i in 1:sN if !any(e->e<0, ss[1:p,i]))
mt = mt[1:p]/mt[end]

#@endexample

# all that's below won't show in doc

@test norm(pathmean(path)-mt)/norm(mt) < 0.1
ess,ns = esspath(path; ns=100)
@test minimum(ess) > 150