text:
function initialize!(integrator,cache::TanYam7ConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 2
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
end

@muladd function perform_step!(integrator, cache::TanYam7ConstantCache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack c1,c2,c3,c4,c5,c6,c7,a21,a31,a32,a41,a43,a51,a53,a54,a61,a63,a64,a65,a71,a73,a74,a75,a76,a81,a83,a84,a85,a86,a87,a91,a93,a94,a95,a96,a97,a98,a101,a103,a104,a105,a106,a107,a108,b1,b4,b5,b6,b7,b8,b9,btilde1,btilde4,btilde5,btilde6,btilde7,btilde8,btilde9,btilde10 = cache
  k1 = integrator.fsalfirst
  a = dt*a21
  k2 = f(uprev+a*k1, p, t + c1*dt)
  k3 = f(uprev+dt*(a31*k1+a32*k2), p, t + c2*dt)
  k4 = f(uprev+dt*(a41*k1       +a43*k3), p, t + c3*dt)
  k5 = f(uprev+dt*(a51*k1       +a53*k3+a54*k4), p, t + c4*dt)
  k6 = f(uprev+dt*(a61*k1       +a63*k3+a64*k4+a65*k5), p, t + c5*dt)
  k7 = f(uprev+dt*(a71*k1       +a73*k3+a74*k4+a75*k5+a76*k6), p, t + c6*dt)
  k8 = f(uprev+dt*(a81*k1       +a83*k3+a84*k4+a85*k5+a86*k6+a87*k7), p, t + c7*dt)
  k9 = f(uprev+dt*(a91*k1       +a93*k3+a94*k4+a95*k5+a96*k6+a97*k7+a98*k8), p, t+dt)
  k10= f(uprev+dt*(a101*k1      +a103*k3+a104*k4+a105*k5+a106*k6+a107*k7+a108*k8), p, t+dt)
  integrator.destats.nf += 9
  u = uprev + dt*(b1*k1+b4*k4+b5*k5+b6*k6+b7*k7+b8*k8+b9*k9)
  if integrator.opts.adaptive
    utilde = dt*(btilde1*k1+btilde4*k4+btilde5*k5+btilde6*k6+btilde7*k7+btilde8*k8+btilde9*k9+btilde10*k10)
    atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
  integrator.fsallast = f(u, p, t+dt) # For the interpolation, needs k at the updated point
  integrator.destats.nf += 1
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.u = u
end

function initialize!(integrator, cache::TanYam7Cache)
  integrator.fsalfirst = cache.fsalfirst
  integrator.fsallast = cache.k
  integrator.kshortsize = 2
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator, cache::TanYam7Cache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack fsalfirst,k2,k3,k4,k5,k6,k7,k8,k9,k10,utilde,tmp,atmp,k = cache
  @unpack c1,c2,c3,c4,c5,c6,c7,a21,a31,a32,a41,a43,a51,a53,a54,a61,a63,a64,a65,a71,a73,a74,a75,a76,a81,a83,a84,a85,a86,a87,a91,a93,a94,a95,a96,a97,a98,a101,a103,a104,a105,a106,a107,a108,b1,b4,b5,b6,b7,b8,b9,btilde1,btilde4,btilde5,btilde6,btilde7,btilde8,btilde9,btilde10 = cache.tab
  k1 = fsalfirst
  f(k1, uprev, p, t)
  a = dt*a21
  @.. tmp = uprev+a*k1
  f(k2, tmp, p, t + c1*dt)
  @.. tmp = uprev+dt*(a31*k1+a32*k2)
  f(k3, tmp, p, t + c2*dt)
  @.. tmp = uprev+dt*(a41*k1+a43*k3)
  f(k4, tmp, p, t + c3*dt)
  @.. tmp = uprev+dt*(a51*k1+a53*k3+a54*k4)
  f(k5, tmp, p, t + c4*dt)
  @.. tmp = uprev+dt*(a61*k1+a63*k3+a64*k4+a65*k5)
  f(k6, tmp, p, t + c5*dt)
  @.. tmp = uprev+dt*(a71*k1+a73*k3+a74*k4+a75*k5+a76*k6)
  f(k7, tmp, p, t + c6*dt)
  @.. tmp = uprev+dt*(a81*k1+a83*k3+a84*k4+a85*k5+a86*k6+a87*k7)
  f(k8, tmp, p, t + c7*dt)
  @.. tmp = uprev+dt*(a91*k1+a93*k3+a94*k4+a95*k5+a96*k6+a97*k7+a98*k8)
  f(k9, tmp, p, t+dt)
  @.. tmp = uprev+dt*(a101*k1+a103*k3+a104*k4+a105*k5+a106*k6+a107*k7+a108*k8)
  f(k10, tmp, p, t+dt)
  @.. u = uprev + dt*(b1*k1+b4*k4+b5*k5+b6*k6+b7*k7+b8*k8+b9*k9)
  integrator.destats.nf += 10
  if integrator.opts.adaptive
    @.. utilde = dt*(btilde1*k1+btilde4*k4+btilde5*k5+btilde6*k6+btilde7*k7+btilde8*k8+btilde9*k9+btilde10*k10)
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
  f(k, u, p, t+dt)
  integrator.destats.nf += 1
  return nothing
end

#=
@muladd function perform_step!(integrator, cache::TanYam7Cache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  uidx = eachindex(integrator.uprev)
  @unpack fsalfirst,k2,k3,k4,k5,k6,k7,k8,k9,k10,utilde,tmp,atmp,k = cache
  @unpack c1,c2,c3,c4,c5,c6,c7,a21,a31,a32,a41,a43,a51,a53,a54,a61,a63,a64,a65,a71,a73,a74,a75,a76,a81,a83,a84,a85,a86,a87,a91,a93,a94,a95,a96,a97,a98,a101,a103,a104,a105,a106,a107,a108,b1,b4,b5,b6,b7,b8,b9,btilde1,btilde4,btilde5,btilde6,btilde7,btilde8,btilde9,btilde10 = cache.tab
  k1 = fsalfirst
  f(k1, uprev, p, t)
  a = dt*a21
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+a*k1[i]
  end
  f(k2, tmp, p, t + c1*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a31*k1[i]+a32*k2[i])
  end
  f(k3, tmp, p, t + c2*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a41*k1[i]+a43*k3[i])
  end
  f(k4, tmp, p, t + c3*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a51*k1[i]+a53*k3[i]+a54*k4[i])
  end
  f(k5, tmp, p, t + c4*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a61*k1[i]+a63*k3[i]+a64*k4[i]+a65*k5[i])
  end
  f(k6, tmp, p, t + c5*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a71*k1[i]+a73*k3[i]+a74*k4[i]+a75*k5[i]+a76*k6[i])
  end
  f(k7, tmp, p, t + c6*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a81*k1[i]+a83*k3[i]+a84*k4[i]+a85*k5[i]+a86*k6[i]+a87*k7[i])
  end
  f(k8, tmp, p, t + c7*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a91*k1[i]+a93*k3[i]+a94*k4[i]+a95*k5[i]+a96*k6[i]+a97*k7[i]+a98*k8[i])
  end
  f(k9, tmp, p, t+dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a101*k1[i]+a103*k3[i]+a104*k4[i]+a105*k5[i]+a106*k6[i]+a107*k7[i]+a108*k8[i])
  end
  f(k10, tmp, p, t+dt)
  @tight_loop_macros for i in uidx
    @inbounds u[i] = uprev[i] + dt*(b1*k1[i]+b4*k4[i]+b5*k5[i]+b6*k6[i]+b7*k7[i]+b8*k8[i]+b9*k9[i])
  end
  integrator.destats.nf += 10
  if integrator.opts.adaptive
    @tight_loop_macros for i in uidx
      @inbounds utilde[i] = dt*(btilde1*k1[i]+btilde4*k4[i]+btilde5*k5[i]+btilde6*k6[i]+btilde7*k7[i]+btilde8*k8[i]+btilde9*k9[i]+btilde10*k10[i])
    end
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
  f(k, u, p, t+dt)
  integrator.destats.nf += 1
end
=#

function initialize!(integrator, cache::DP8ConstantCache)
  integrator.kshortsize = 7
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  @inbounds for i in eachindex(integrator.k)
    integrator.k[i] = zero(integrator.fsalfirst)
  end
end

@muladd function perform_step!(integrator, cache::DP8ConstantCache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack c7,c8,c9,c10,c11,c6,c5,c4,c3,c2,b1,b6,b7,b8,b9,b10,b11,b12,btilde1,btilde6,btilde7,btilde8,btilde9,btilde10,btilde11,btilde12,er1,er6,er7,er8,er9,er10,er11,er12,a0201,a0301,a0302,a0401,a0403,a0501,a0503,a0504,a0601,a0604,a0605,a0701,a0704,a0705,a0706,a0801,a0804,a0805,a0806,a0807,a0901,a0904,a0905,a0906,a0907,a0908,a1001,a1004,a1005,a1006,a1007,a1008,a1009,a1101,a1104,a1105,a1106,a1107,a1108,a1109,a1110,a1201,a1204,a1205,a1206,a1207,a1208,a1209,a1210,a1211 = cache
  k1 = integrator.fsalfirst
  a = dt*a0201
  k2 = f(uprev+a*k1, p, t + c2*dt)
  k3 = f(uprev+dt*(a0301*k1+a0302*k2), p, t + c3*dt)
  k4 = f(uprev+dt*(a0401*k1       +a0403*k3), p, t + c4*dt)
  k5 = f(uprev+dt*(a0501*k1       +a0503*k3+a0504*k4), p, t + c5*dt)
  k6 = f(uprev+dt*(a0601*k1                +a0604*k4+a0605*k5), p, t + c6*dt)
  k7 = f(uprev+dt*(a0701*k1                +a0704*k4+a0705*k5+a0706*k6), p, t + c7*dt)
  k8 = f(uprev+dt*(a0801*k1                +a0804*k4+a0805*k5+a0806*k6+a0807*k7), p, t + c8*dt)
  k9 = f(uprev+dt*(a0901*k1                +a0904*k4+a0905*k5+a0906*k6+a0907*k7+a0908*k8), p, t + c9*dt)
  k10 =f(uprev+dt*(a1001*k1                +a1004*k4+a1005*k5+a1006*k6+a1007*k7+a1008*k8+a1009*k9), p, t + c10*dt)
  k11= f(uprev+dt*(a1101*k1                +a1104*k4+a1105*k5+a1106*k6+a1107*k7+a1108*k8+a1109*k9+a1110*k10), p, t + c11*dt)
  k12= f(uprev+dt*(a1201*k1                +a1204*k4+a1205*k5+a1206*k6+a1207*k7+a1208*k8+a1209*k9+a1210*k10+a1211*k11), p, t+dt)
  integrator.destats.nf += 11
  kupdate= b1*k1+b6*k6+b7*k7+b8*k8+b9*k9+b10*k10+b11*k11+b12*k12
  u = uprev + dt*kupdate
  if integrator.opts.adaptive
    utilde = dt*(k1*er1 + k6*er6 + k7*er7 + k8*er8 + k9*er9 + k10*er10 + k11*er11 + k12*er12)
    atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    err5 = integrator.opts.internalnorm(atmp,t) # Order 5
    utilde = dt*(btilde1*k1 + btilde6*k6 + btilde7*k7 + btilde8*k8 + btilde9*k9 + btilde10*k10 + btilde11*k11 + btilde12*k12)
    atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    err3 = integrator.opts.internalnorm(atmp,t) # Order 3
    err52 = err5*err5
    if err5 ≈ 0 && err3 ≈ 0
      integrator.EEst = zero(integrator.EEst)
    else
      integrator.EEst = err52/sqrt(err52 + 0.01*err3*err3)
    end
  end
  k13 = f(u, p, t+dt)
  integrator.destats.nf += 1
  integrator.fsallast = k13
  if integrator.opts.calck
    @unpack c14,c15,c16,a1401,a1407,a1408,a1409,a1410,a1411,a1412,a1413,a1501,a1506,a1507,a1508,a1511,a1512,a1513,a1514,a1601,a1606,a1607,a1608,a1609,a1613,a1614,a1615 = cache
    @unpack d401,d406,d407,d408,d409,d410,d411,d412,d413,d414,d415,d416,d501,d506,d507,d508,d509,d510,d511,d512,d513,d514,d515,d516,d601,d606,d607,d608,d609,d610,d611,d612,d613,d614,d615,d616,d701,d706,d707,d708,d709,d710,d711,d712,d713,d714,d715,d716 = cache
    k14 = f(uprev+dt*(a1401*k1         +a1407*k7+a1408*k8+a1409*k9+a1410*k10+a1411*k11+a1412*k12+a1413*k13), p, t + c14*dt)
    k15 = f(uprev+dt*(a1501*k1+a1506*k6+a1507*k7+a1508*k8                   +a1511*k11+a1512*k12+a1513*k13+a1514*k14), p, t + c15*dt)
    k16 = f(uprev+dt*(a1601*k1+a1606*k6+a1607*k7+a1608*k8+a1609*k9                              +a1613*k13+a1614*k14+a1615*k15), p, t + c16*dt)
    integrator.destats.nf += 3
    udiff = kupdate
    integrator.k[1] = udiff
    bspl = k1 - udiff
    integrator.k[2] = bspl
    integrator.k[3] = udiff - k13 - bspl
    integrator.k[4] = d401*k1+d406*k6+d407*k7+d408*k8+d409*k9+d410*k10+d411*k11+d412*k12+d413*k13+d414*k14+d415*k15+d416*k16
    integrator.k[5] = d501*k1+d506*k6+d507*k7+d508*k8+d509*k9+d510*k10+d511*k11+d512*k12+d513*k13+d514*k14+d515*k15+d516*k16
    integrator.k[6] = d601*k1+d606*k6+d607*k7+d608*k8+d609*k9+d610*k10+d611*k11+d612*k12+d613*k13+d614*k14+d615*k15+d616*k16
    integrator.k[7] = d701*k1+d706*k6+d707*k7+d708*k8+d709*k9+d710*k10+d711*k11+d712*k12+d713*k13+d714*k14+d715*k15+d716*k16
  end
  integrator.u = u
end

function initialize!(integrator, cache::DP8Cache)
  integrator.kshortsize = 7
  resize!(integrator.k, integrator.kshortsize)
  integrator.k .= [cache.udiff,cache.bspl,cache.dense_tmp3,cache.dense_tmp4,cache.dense_tmp5,cache.dense_tmp6,cache.dense_tmp7]
  integrator.fsalfirst = cache.k1
  integrator.fsallast = cache.k13
  integrator.f(integrator.fsalfirst,integrator.uprev,integrator.p,integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator, cache::DP8Cache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  uidx = eachindex(integrator.uprev)
  @unpack c7,c8,c9,c10,c11,c6,c5,c4,c3,c2,b1,b6,b7,b8,b9,b10,b11,b12,btilde1,btilde6,btilde7,btilde8,btilde9,btilde10,btilde11,btilde12,er1,er6,er7,er8,er9,er10,er11,er12,a0201,a0301,a0302,a0401,a0403,a0501,a0503,a0504,a0601,a0604,a0605,a0701,a0704,a0705,a0706,a0801,a0804,a0805,a0806,a0807,a0901,a0904,a0905,a0906,a0907,a0908,a1001,a1004,a1005,a1006,a1007,a1008,a1009,a1101,a1104,a1105,a1106,a1107,a1108,a1109,a1110,a1201,a1204,a1205,a1206,a1207,a1208,a1209,a1210,a1211 = cache.tab
  @unpack k1,k2,k3,k4,k5,k6,k7,k8,k9,k10,k11,k12,k13,k14,k15,k16,udiff,bspl,dense_tmp3,dense_tmp4,dense_tmp5,dense_tmp6,dense_tmp7,kupdate,utilde,tmp,atmp = cache
  f(k1, uprev, p, t)
  a = dt*a0201
  @.. tmp = uprev+a*k1
  f(k2, tmp, p, t + c2*dt)
  @.. tmp = uprev+dt*(a0301*k1+a0302*k2)
  f(k3, tmp, p, t + c3*dt)
  @.. tmp = uprev+dt*(a0401*k1+a0403*k3)
  f(k4, tmp, p, t + c4*dt)
  @.. tmp = uprev+dt*(a0501*k1+a0503*k3+a0504*k4)
  f(k5, tmp, p, t + c5*dt)
  @.. tmp = uprev+dt*(a0601*k1+a0604*k4+a0605*k5)
  f(k6, tmp, p, t + c6*dt)
  @.. tmp = uprev+dt*(a0701*k1+a0704*k4+a0705*k5+a0706*k6)
  f(k7, tmp, p, t + c7*dt)
  @.. tmp = uprev+dt*(a0801*k1+a0804*k4+a0805*k5+a0806*k6+a0807*k7)
  f(k8, tmp, p, t + c8*dt)
  @.. tmp = uprev+dt*(a0901*k1+a0904*k4+a0905*k5+a0906*k6+a0907*k7+a0908*k8)
  f(k9, tmp, p, t + c9*dt)
  @.. tmp = uprev+dt*(a1001*k1+a1004*k4+a1005*k5+a1006*k6+a1007*k7+a1008*k8+a1009*k9)
  f(k10, tmp, p, t + c10*dt)
  @.. tmp = uprev+dt*(a1101*k1+a1104*k4+a1105*k5+a1106*k6+a1107*k7+a1108*k8+a1109*k9+a1110*k10)
  f(k11, tmp, p, t + c11*dt)
  @.. tmp = uprev+dt*(a1201*k1+a1204*k4+a1205*k5+a1206*k6+a1207*k7+a1208*k8+a1209*k9+a1210*k10+a1211*k11)
  f(k12, tmp, p, t+dt)
  @.. kupdate = b1*k1+b6*k6+b7*k7+b8*k8+b9*k9+b10*k10+b11*k11+b12*k12
  @.. u = uprev + dt*kupdate
  integrator.destats.nf += 12
  if integrator.opts.adaptive
    @.. utilde = dt*(k1*er1 + k6*er6 + k7*er7 + k8*er8 + k9*er9 + k10*er10 + k11*er11 + k12*er12)
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    err5 = integrator.opts.internalnorm(atmp,t) # Order 5
    @.. utilde= dt*(btilde1*k1 + btilde6*k6 + btilde7*k7 + btilde8*k8 + btilde9*k9 + btilde10*k10 + btilde11*k11 + btilde12*k12)
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    err3 = integrator.opts.internalnorm(atmp,t) # Order 3
    err52 = err5*err5
    if err5 ≈ 0 && err3 ≈ 0
      integrator.EEst = zero(integrator.EEst)
    else
      integrator.EEst = err52/sqrt(err52 + 0.01*err3*err3)
    end
  end
  f(k13, u, p, t+dt)
  integrator.destats.nf += 1
  if integrator.opts.calck
    @unpack c14,c15,c16,a1401,a1407,a1408,a1409,a1410,a1411,a1412,a1413,a1501,a1506,a1507,a1508,a1511,a1512,a1513,a1514,a1601,a1606,a1607,a1608,a1609,a1613,a1614,a1615 = cache.tab
    @unpack d401,d406,d407,d408,d409,d410,d411,d412,d413,d414,d415,d416,d501,d506,d507,d508,d509,d510,d511,d512,d513,d514,d515,d516,d601,d606,d607,d608,d609,d610,d611,d612,d613,d614,d615,d616,d701,d706,d707,d708,d709,d710,d711,d712,d713,d714,d715,d716 = cache.tab
    @.. tmp = uprev+dt*(a1401*k1+a1407*k7+a1408*k8+a1409*k9+a1410*k10+a1411*k11+a1412*k12+a1413*k13)
    f(k14, tmp, p, t + c14*dt)
    @.. tmp = uprev+dt*(a1501*k1+a1506*k6+a1507*k7+a1508*k8+a1511*k11+a1512*k12+a1513*k13+a1514*k14)
    f(k15, tmp, p, t + c15*dt)
    @.. tmp = uprev+dt*(a1601*k1+a1606*k6+a1607*k7+a1608*k8+a1609*k9+a1613*k13+a1614*k14+a1615*k15)
    f(k16, tmp, p, t + c16*dt)
    integrator.destats.nf += 3
    @.. udiff= kupdate
    @.. bspl = k1 - udiff
    @.. integrator.k[3] = udiff - k13 - bspl
    @.. integrator.k[4] = d401*k1+d406*k6+d407*k7+d408*k8+d409*k9+d410*k10+d411*k11+d412*k12+d413*k13+d414*k14+d415*k15+d416*k16
    @.. integrator.k[5] = d501*k1+d506*k6+d507*k7+d508*k8+d509*k9+d510*k10+d511*k11+d512*k12+d513*k13+d514*k14+d515*k15+d516*k16
    @.. integrator.k[6] = d601*k1+d606*k6+d607*k7+d608*k8+d609*k9+d610*k10+d611*k11+d612*k12+d613*k13+d614*k14+d615*k15+d616*k16
    @.. integrator.k[7] = d701*k1+d706*k6+d707*k7+d708*k8+d709*k9+d710*k10+d711*k11+d712*k12+d713*k13+d714*k14+d715*k15+d716*k16
  end
end

#=
@muladd function perform_step!(integrator, cache::DP8Cache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  uidx = eachindex(integrator.uprev)
  @unpack c7,c8,c9,c10,c11,c6,c5,c4,c3,c2,b1,b6,b7,b8,b9,b10,b11,b12,btilde1,btilde6,btilde7,btilde8,btilde9,btilde10,btilde11,btilde12,er1,er6,er7,er8,er9,er10,er11,er12,a0201,a0301,a0302,a0401,a0403,a0501,a0503,a0504,a0601,a0604,a0605,a0701,a0704,a0705,a0706,a0801,a0804,a0805,a0806,a0807,a0901,a0904,a0905,a0906,a0907,a0908,a1001,a1004,a1005,a1006,a1007,a1008,a1009,a1101,a1104,a1105,a1106,a1107,a1108,a1109,a1110,a1201,a1204,a1205,a1206,a1207,a1208,a1209,a1210,a1211 = cache.tab
  @unpack k1,k2,k3,k4,k5,k6,k7,k8,k9,k10,k11,k12,k13,k14,k15,k16,udiff,bspl,dense_tmp3,dense_tmp4,dense_tmp5,dense_tmp6,dense_tmp7,kupdate,utilde,tmp,atmp = cache
  f(k1, uprev, p, t)
  a = dt*a0201
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+a*k1[i]
  end
  f(k2, tmp, p, t + c2*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0301*k1[i]+a0302*k2[i])
  end
  f(k3, tmp, p, t + c3*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0401*k1[i]+a0403*k3[i])
  end
  f(k4, tmp, p, t + c4*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0501*k1[i]+a0503*k3[i]+a0504*k4[i])
  end
  f(k5, tmp, p, t + c5*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0601*k1[i]+a0604*k4[i]+a0605*k5[i])
  end
  f(k6, tmp, p, t + c6*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0701*k1[i]+a0704*k4[i]+a0705*k5[i]+a0706*k6[i])
  end
  f(k7, tmp, p, t + c7*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0801*k1[i]+a0804*k4[i]+a0805*k5[i]+a0806*k6[i]+a0807*k7[i])
  end
  f(k8, tmp, p, t + c8*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0901*k1[i]+a0904*k4[i]+a0905*k5[i]+a0906*k6[i]+a0907*k7[i]+a0908*k8[i])
  end
  f(k9, tmp, p, t + c9*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a1001*k1[i]+a1004*k4[i]+a1005*k5[i]+a1006*k6[i]+a1007*k7[i]+a1008*k8[i]+a1009*k9[i])
  end
  f(k10, tmp, p, t + c10*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a1101*k1[i]+a1104*k4[i]+a1105*k5[i]+a1106*k6[i]+a1107*k7[i]+a1108*k8[i]+a1109*k9[i]+a1110*k10[i])
  end
  f(k11, tmp, p, t + c11*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a1201*k1[i]+a1204*k4[i]+a1205*k5[i]+a1206*k6[i]+a1207*k7[i]+a1208*k8[i]+a1209*k9[i]+a1210*k10[i]+a1211*k11[i])
  end
  f(k12, tmp, p, t+dt)
  @tight_loop_macros for i in uidx
    @inbounds kupdate[i] = b1*k1[i]+b6*k6[i]+b7*k7[i]+b8*k8[i]+b9*k9[i]+b10*k10[i]+b11*k11[i]+b12*k12[i]
    @inbounds u[i] = uprev[i] + dt*kupdate[i]
  end
  integrator.destats.nf += 12
  if integrator.opts.adaptive
    @tight_loop_macros for i in uidx
      @inbounds utilde[i] = dt*(k1[i]*er1 + k6[i]*er6 + k7[i]*er7 + k8[i]*er8 + k9[i]*er9 + k10[i]*er10 + k11[i]*er11 + k12[i]*er12)
    end
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    err5 = integrator.opts.internalnorm(atmp,t) # Order 5
    @tight_loop_macros for i in uidx
      @inbounds utilde[i]= dt*(btilde1*k1[i] + btilde6*k6[i] + btilde7*k7[i] + btilde8*k8[i] + btilde9*k9[i] + btilde10*k10[i] + btilde11*k11[i] + btilde12*k12[i])
    end
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    err3 = integrator.opts.internalnorm(atmp,t) # Order 3
    err52 = err5*err5
    if err5 ≈ 0 && err3 ≈ 0
      integrator.EEst = zero(integrator.EEst)
    else
      integrator.EEst = err52/sqrt(err52 + 0.01*err3*err3)
    end
  end
  f(k13, u, p, t+dt)
  integrator.destats.nf += 1
  if integrator.opts.calck
    @unpack c14,c15,c16,a1401,a1407,a1408,a1409,a1410,a1411,a1412,a1413,a1501,a1506,a1507,a1508,a1511,a1512,a1513,a1514,a1601,a1606,a1607,a1608,a1609,a1613,a1614,a1615 = cache.tab
    @unpack d401,d406,d407,d408,d409,d410,d411,d412,d413,d414,d415,d416,d501,d506,d507,d508,d509,d510,d511,d512,d513,d514,d515,d516,d601,d606,d607,d608,d609,d610,d611,d612,d613,d614,d615,d616,d701,d706,d707,d708,d709,d710,d711,d712,d713,d714,d715,d716 = cache.tab
    @tight_loop_macros for i in uidx
      @inbounds tmp[i] = uprev[i]+dt*(a1401*k1[i]+a1407*k7[i]+a1408*k8[i]+a1409*k9[i]+a1410*k10[i]+a1411*k11[i]+a1412*k12[i]+a1413*k13[i])
    end
    f(k14, tmp, p, t + c14*dt)
    @tight_loop_macros for i in uidx
      @inbounds tmp[i] = uprev[i]+dt*(a1501*k1[i]+a1506*k6[i]+a1507*k7[i]+a1508*k8[i]+a1511*k11[i]+a1512*k12[i]+a1513*k13[i]+a1514*k14[i])
    end
    f(k15, tmp, p, t + c15*dt)
    @tight_loop_macros for i in uidx
      @inbounds tmp[i] = uprev[i]+dt*(a1601*k1[i]+a1606*k6[i]+a1607*k7[i]+a1608*k8[i]+a1609*k9[i]+a1613*k13[i]+a1614*k14[i]+a1615*k15[i])
    end
    f(k16, tmp, p, t + c16*dt)
    integrator.destats.nf += 3
    @tight_loop_macros for i in uidx
      @inbounds udiff[i]= kupdate[i]
      @inbounds bspl[i] = k1[i] - udiff[i]
      @inbounds integrator.k[3][i] = udiff[i] - k13[i] - bspl[i]
      @inbounds integrator.k[4][i] = d401*k1[i]+d406*k6[i]+d407*k7[i]+d408*k8[i]+d409*k9[i]+d410*k10[i]+d411*k11[i]+d412*k12[i]+d413*k13[i]+d414*k14[i]+d415*k15[i]+d416*k16[i]
      @inbounds integrator.k[5][i] = d501*k1[i]+d506*k6[i]+d507*k7[i]+d508*k8[i]+d509*k9[i]+d510*k10[i]+d511*k11[i]+d512*k12[i]+d513*k13[i]+d514*k14[i]+d515*k15[i]+d516*k16[i]
      @inbounds integrator.k[6][i] = d601*k1[i]+d606*k6[i]+d607*k7[i]+d608*k8[i]+d609*k9[i]+d610*k10[i]+d611*k11[i]+d612*k12[i]+d613*k13[i]+d614*k14[i]+d615*k15[i]+d616*k16[i]
      @inbounds integrator.k[7][i] = d701*k1[i]+d706*k6[i]+d707*k7[i]+d708*k8[i]+d709*k9[i]+d710*k10[i]+d711*k11[i]+d712*k12[i]+d713*k13[i]+d714*k14[i]+d715*k15[i]+d716*k16[i]
    end
  end
end
=#

function initialize!(integrator, cache::TsitPap8ConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 2
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
end

#=
@muladd function perform_step!(integrator, cache::TsitPap8ConstantCache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,a0201,a0301,a0302,a0401,a0403,a0501,a0503,a0504,a0601,a0604,a0605,a0701,a0704,a0705,a0706,a0801,a0804,a0805,a0806,a0807,a0901,a0904,a0905,a0906,a0907,a0908,a1001,a1004,a1005,a1006,a1007,a1008,a1009,a1101,a1104,a1105,a1106,a1107,a1108,a1109,a1110,a1201,a1204,a1205,a1206,a1207,a1208,a1209,a1210,a1211,a1301,a1304,a1305,a1306,a1307,a1308,a1309,a1310,b1,b6,b7,b8,b9,b10,b11,b12,btilde1,btilde6,btilde7,btilde8,btilde9,btilde10,btilde11,btilde12,btilde13 = cache
  k1 = integrator.fsalfirst
  a = dt*a0201
  k2 = f(t + c1*dt, @.. uprev+a*k1)
  k3 = f(t + c2*dt, @.. uprev+dt*(a0301*k1+a0302*k2))
  k4 = f(t + c3*dt, @.. uprev+dt*(a0401*k1       +a0403*k3))
  k5 = f(t + c4*dt, @.. uprev+dt*(a0501*k1       +a0503*k3+a0504*k4))
  k6 = f(t + c5*dt, @.. uprev+dt*(a0601*k1                +a0604*k4+a0605*k5))
  k7 = f(t + c6*dt, @.. uprev+dt*(a0701*k1                +a0704*k4+a0705*k5+a0706*k6))
  k8 = f(t + c7*dt, @.. uprev+dt*(a0801*k1                +a0804*k4+a0805*k5+a0806*k6+a0807*k7))
  k9 = f(t + c8*dt, @.. uprev+dt*(a0901*k1                +a0904*k4+a0905*k5+a0906*k6+a0907*k7+a0908*k8))
  k10 =f(t + c9*dt, @.. uprev+dt*(a1001*k1                +a1004*k4+a1005*k5+a1006*k6+a1007*k7+a1008*k8+a1009*k9))
  k11= f(t + c10*dt, @.. uprev+dt*(a1101*k1                +a1104*k4+a1105*k5+a1106*k6+a1107*k7+a1108*k8+a1109*k9+a1110*k10))
  k12= f(t+dt, @.. uprev+dt*(a1201*k1                +a1204*k4+a1205*k5+a1206*k6+a1207*k7+a1208*k8+a1209*k9+a1210*k10+a1211*k11))
  k13= f(t+dt, @.. uprev+dt*(a1301*k1                +a1304*k4+a1305*k5+a1306*k6+a1307*k7+a1308*k8+a1309*k9+a1310*k10))
  u = @.. uprev + dt*(b1*k1+b6*k6+b7*k7+b8*k8+b9*k9+b10*k10+b11*k11+b12*k12)
  if integrator.opts.adaptive
    utilde = @.. dt*(btilde1*k1 + btilde6*k6 + btilde7*k7 + btilde8*k8 + btilde9*k9 + btilde10*k10 + btilde11*k11 + btilde12*k12 + btilde13*k13)
    atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
  integrator.fsallast = f(u, p, t+dt)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.u = u
end
=#

@muladd function perform_step!(integrator, cache::TsitPap8ConstantCache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,a0201,a0301,a0302,a0401,a0403,a0501,a0503,a0504,a0601,a0604,a0605,a0701,a0704,a0705,a0706,a0801,a0804,a0805,a0806,a0807,a0901,a0904,a0905,a0906,a0907,a0908,a1001,a1004,a1005,a1006,a1007,a1008,a1009,a1101,a1104,a1105,a1106,a1107,a1108,a1109,a1110,a1201,a1204,a1205,a1206,a1207,a1208,a1209,a1210,a1211,a1301,a1304,a1305,a1306,a1307,a1308,a1309,a1310,b1,b6,b7,b8,b9,b10,b11,b12,btilde1,btilde6,btilde7,btilde8,btilde9,btilde10,btilde11,btilde12,btilde13 = cache
  k1 = integrator.fsalfirst
  a = dt*a0201
  k2 = f(uprev+a*k1, p, t + c1*dt)
  k3 = f(uprev+dt*(a0301*k1+a0302*k2), p, t + c2*dt)
  k4 = f(uprev+dt*(a0401*k1       +a0403*k3), p, t + c3*dt)
  k5 = f(uprev+dt*(a0501*k1       +a0503*k3+a0504*k4), p, t + c4*dt)
  k6 = f(uprev+dt*(a0601*k1                +a0604*k4+a0605*k5), p, t + c5*dt)
  k7 = f(uprev+dt*(a0701*k1                +a0704*k4+a0705*k5+a0706*k6), p, t + c6*dt)
  k8 = f(uprev+dt*(a0801*k1                +a0804*k4+a0805*k5+a0806*k6+a0807*k7), p, t + c7*dt)
  k9 = f(uprev+dt*(a0901*k1                +a0904*k4+a0905*k5+a0906*k6+a0907*k7+a0908*k8), p, t + c8*dt)
  k10 =f(uprev+dt*(a1001*k1                +a1004*k4+a1005*k5+a1006*k6+a1007*k7+a1008*k8+a1009*k9), p, t + c9*dt)
  k11= f(uprev+dt*(a1101*k1                +a1104*k4+a1105*k5+a1106*k6+a1107*k7+a1108*k8+a1109*k9+a1110*k10), p, t + c10*dt)
  k12= f(uprev+dt*(a1201*k1                +a1204*k4+a1205*k5+a1206*k6+a1207*k7+a1208*k8+a1209*k9+a1210*k10+a1211*k11), p, t+dt)
  k13= f(uprev+dt*(a1301*k1                +a1304*k4+a1305*k5+a1306*k6+a1307*k7+a1308*k8+a1309*k9+a1310*k10), p, t+dt)
  integrator.destats.nf += 12
  u = uprev + dt*(b1*k1+b6*k6+b7*k7+b8*k8+b9*k9+b10*k10+b11*k11+b12*k12)
  if integrator.opts.adaptive
    utilde = dt*(btilde1*k1 + btilde6*k6 + btilde7*k7 + btilde8*k8 + btilde9*k9 + btilde10*k10 + btilde11*k11 + btilde12*k12 + btilde13*k13)
    atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
  integrator.fsallast = f(u, p, t+dt)
  integrator.destats.nf += 1
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.u = u
end

function initialize!(integrator, cache::TsitPap8Cache)
  integrator.fsalfirst = cache.fsalfirst
  integrator.fsallast = cache.k
  integrator.kshortsize = 2
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator, cache::TsitPap8Cache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,a0201,a0301,a0302,a0401,a0403,a0501,a0503,a0504,a0601,a0604,a0605,a0701,a0704,a0705,a0706,a0801,a0804,a0805,a0806,a0807,a0901,a0904,a0905,a0906,a0907,a0908,a1001,a1004,a1005,a1006,a1007,a1008,a1009,a1101,a1104,a1105,a1106,a1107,a1108,a1109,a1110,a1201,a1204,a1205,a1206,a1207,a1208,a1209,a1210,a1211,a1301,a1304,a1305,a1306,a1307,a1308,a1309,a1310,b1,b6,b7,b8,b9,b10,b11,b12,btilde1,btilde6,btilde7,btilde8,btilde9,btilde10,btilde11,btilde12,btilde13 = cache.tab
  @unpack k2,k3,k4,k5,k6,k7,k8,k9,k10,k11,k12,k13,utilde,tmp,atmp,k = cache
  k1 = cache.fsalfirst
  f(k1, uprev, p, t)
  a = dt*a0201
  @.. tmp = uprev+a*k1
  f(k2, tmp, p, t + c1*dt)
  @.. tmp = uprev+dt*(a0301*k1+a0302*k2)
  f(k3, tmp, p, t + c2*dt)
  @.. tmp = uprev+dt*(a0401*k1+a0403*k3)
  f(k4, tmp, p, t + c3*dt)
  @.. tmp = uprev+dt*(a0501*k1+a0503*k3+a0504*k4)
  f(k5, tmp, p, t + c4*dt)
  @.. tmp = uprev+dt*(a0601*k1+a0604*k4+a0605*k5)
  f(k6, tmp, p, t + c5*dt)
  @.. tmp = uprev+dt*(a0701*k1+a0704*k4+a0705*k5+a0706*k6)
  f(k7, tmp, p, t + c6*dt)
  @.. tmp = uprev+dt*(a0801*k1+a0804*k4+a0805*k5+a0806*k6+a0807*k7)
  f(k8, tmp, p, t + c7*dt)
  @.. tmp = uprev+dt*(a0901*k1+a0904*k4+a0905*k5+a0906*k6+a0907*k7+a0908*k8)
  f(k9, tmp, p, t + c8*dt)
  @.. tmp = uprev+dt*(a1001*k1+a1004*k4+a1005*k5+a1006*k6+a1007*k7+a1008*k8+a1009*k9)
  f(k10, tmp, p, t + c9*dt)
  @.. tmp = uprev+dt*(a1101*k1+a1104*k4+a1105*k5+a1106*k6+a1107*k7+a1108*k8+a1109*k9+a1110*k10)
  f(k11, tmp, p, t + c10*dt)
  @.. tmp = uprev+dt*(a1201*k1+a1204*k4+a1205*k5+a1206*k6+a1207*k7+a1208*k8+a1209*k9+a1210*k10+a1211*k11)
  f(k12, tmp, p, t+dt)
  @.. tmp = uprev+dt*(a1301*k1+a1304*k4+a1305*k5+a1306*k6+a1307*k7+a1308*k8+a1309*k9+a1310*k10)
  f(k13, tmp, p, t+dt)
  @.. u = uprev + dt*(b1*k1+b6*k6+b7*k7+b8*k8+b9*k9+b10*k10+b11*k11+b12*k12)
  integrator.destats.nf += 13
  if integrator.opts.adaptive
    @.. utilde = dt*(btilde1*k1 + btilde6*k6 + btilde7*k7 + btilde8*k8 + btilde9*k9 + btilde10*k10 + btilde11*k11 + btilde12*k12 + btilde13*k13)
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
  f(k, u, p, t+dt)
  integrator.destats.nf += 1
  return nothing
end

#=
@muladd function perform_step!(integrator, cache::TsitPap8Cache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  uidx = eachindex(integrator.uprev)
  @unpack c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,a0201,a0301,a0302,a0401,a0403,a0501,a0503,a0504,a0601,a0604,a0605,a0701,a0704,a0705,a0706,a0801,a0804,a0805,a0806,a0807,a0901,a0904,a0905,a0906,a0907,a0908,a1001,a1004,a1005,a1006,a1007,a1008,a1009,a1101,a1104,a1105,a1106,a1107,a1108,a1109,a1110,a1201,a1204,a1205,a1206,a1207,a1208,a1209,a1210,a1211,a1301,a1304,a1305,a1306,a1307,a1308,a1309,a1310,b1,b6,b7,b8,b9,b10,b11,b12,btilde1,btilde6,btilde7,btilde8,btilde9,btilde10,btilde11,btilde12,btilde13 = cache.tab
  @unpack k2,k3,k4,k5,k6,k7,k8,k9,k10,k11,k12,k13,utilde,tmp,atmp,k = cache
  k1 = cache.fsalfirst
  f(k1, uprev, p, t)
  a = dt*a0201
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+a*k1[i]
  end
  f(k2, tmp, p, t + c1*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0301*k1[i]+a0302*k2[i])
  end
  f(k3, tmp, p, t + c2*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0401*k1[i]+a0403*k3[i])
  end
  f(k4, tmp, p, t + c3*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0501*k1[i]+a0503*k3[i]+a0504*k4[i])
  end
  f(k5, tmp, p, t + c4*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0601*k1[i]+a0604*k4[i]+a0605*k5[i])
  end
  f(k6, tmp, p, t + c5*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0701*k1[i]+a0704*k4[i]+a0705*k5[i]+a0706*k6[i])
  end
  f(k7, tmp, p, t + c6*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0801*k1[i]+a0804*k4[i]+a0805*k5[i]+a0806*k6[i]+a0807*k7[i])
  end
  f(k8, tmp, p, t + c7*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a0901*k1[i]+a0904*k4[i]+a0905*k5[i]+a0906*k6[i]+a0907*k7[i]+a0908*k8[i])
  end
  f(k9, tmp, p, t + c8*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a1001*k1[i]+a1004*k4[i]+a1005*k5[i]+a1006*k6[i]+a1007*k7[i]+a1008*k8[i]+a1009*k9[i])
  end
  f(k10, tmp, p, t + c9*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a1101*k1[i]+a1104*k4[i]+a1105*k5[i]+a1106*k6[i]+a1107*k7[i]+a1108*k8[i]+a1109*k9[i]+a1110*k10[i])
  end
  f(k11, tmp, p, t + c10*dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a1201*k1[i]+a1204*k4[i]+a1205*k5[i]+a1206*k6[i]+a1207*k7[i]+a1208*k8[i]+a1209*k9[i]+a1210*k10[i]+a1211*k11[i])
  end
  f(k12, tmp, p, t+dt)
  @tight_loop_macros for i in uidx
    @inbounds tmp[i] = uprev[i]+dt*(a1301*k1[i]+a1304*k4[i]+a1305*k5[i]+a1306*k6[i]+a1307*k7[i]+a1308*k8[i]+a1309*k9[i]+a1310*k10[i])
  end
  f(k13, tmp, p, t+dt)
  @tight_loop_macros for i in uidx
    @inbounds u[i] = uprev[i] + dt*(b1*k1[i]+b6*k6[i]+b7*k7[i]+b8*k8[i]+b9*k9[i]+b10*k10[i]+b11*k11[i]+b12*k12[i])
  end
  integrator.destats.nf += 13
  if integrator.opts.adaptive
    @tight_loop_macros for i in uidx
      @inbounds utilde[i] = dt*(btilde1*k1[i] + btilde6*k6[i] + btilde7*k7[i] + btilde8*k8[i] + btilde9*k9[i] + btilde10*k10[i] + btilde11*k11[i] + btilde12*k12[i] + btilde13*k13[i])
    end
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
  f(k, u, p, t+dt)
  integrator.destats.nf += 1
end
=#

function initialize!(integrator, cache::PFRK87ConstantCache)
  integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
  integrator.kshortsize = 2
  integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)

  # Avoid undefined entries if k is an array of arrays
  integrator.fsallast = zero(integrator.fsalfirst)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
end

@muladd function perform_step!(integrator, cache::PFRK87ConstantCache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack α0201, α0301, α0401, α0501, α0601, α0701, α0302, α0403, α0503, α0504, α0604, α0704, α0605, α0705, α0706, α0908, α1008, α1108, α1208, α1308, α1009, α1109, α1209, α1309, α1110, α1210, α1310, α1211, α1311, β1, β6, β7, β8, β9, β10, β11, β12, β13, β1tilde, β6tilde, β7tilde, β8tilde, β9tilde, β10tilde, β11tilde, β12tilde, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13 = cache
  alg = unwrap_alg(integrator, true)
  ν = alg.omega*dt
  νsq = ν^2

  c_ν = -0.19781108078634084 - (0.164050909125528499*νsq) + (0.042578310088756321*(νsq^2)) - (0.002300513610963998*(νsq^3)) + (0.000033467244551879287*(νsq^4)) - (7.8661142036921924*(1/100000000)*(νsq^5))
  d_ν = 1-(0.296457092123567400*νsq) + (0.0015793885907465726*(νsq^2)) - (0.00018913011771688527*(νsq^3)) + (0.000017089234650765179*(νsq^4)) - (1.2705211682518626*(1/10000000)*(νsq^5))

  α0807 = c_ν/d_ν
  α0801 = 0.026876256 + 0.0576576*α0807
  α0804 = 0.22464336 - 0.944944*α0807
  α0805 = 0.000369024 - 0.2061696*α0807
  α0806 = 0.21311136 + 0.093456*α0807
  α0901 = 0.07239997637512857 + 0.01913119863380767*α0807
  α0904 = -0.688400520601143 - 0.3135390887207368*α0807
  α0905 = -0.17301267570583073 - 0.06840852844816077*α0807
  α0906 = 0.1440060555560846 + 0.031009360422930017*α0807
  α0907 = 0.9982362892760762 + 0.33180705811215994*α0807
  α1001 = 0.16261514523236525 - 0.12125171966747463*α0807
  α1004 = -2.1255544052061124 + 1.9871809612169453*α0807
  α1005 = -0.216403903283323 + 0.43356675517460624*α0807
  α1006 = -0.060417230254934076 - 0.1965343807796979*α0807
  α1007 = 2.4846281621788395 - 2.102961615944379*α0807
  α1101 = -1.0320124180911034 + 1.061943768952537*α0807
  α1104 = 13.666683232895137 - 17.40407843561103*α0807
  α1105 = 0.25990355211486116 - 3.797253476860588*α0807
  α1106 = -5.759316475814002 + 1.7212824826428488*α0807
  α1107 = -12.822511612651839 + 18.41810566087623*α0807
  α1201 = 0.2478349764611783 - 0.06383934946543009*α0807
  α1204 = -4.593782880309185 + 1.046256005127882*α0807
  α1205 = -0.39566692537411896 + 0.22827403748244698*α0807
  α1206 = -3.0673550479691665 - 0.10347586863902129*α0807
  α1207 = 5.386688702227177 - 1.1072148245058775*α0807
  α1301 = 0.7332242174431163 - 0.5164807626867616*α0807
  α1304 = -10.196728938160977 + 8.464545832921925*α0807
  α1305 = -0.43865244706547707 + 1.846809999910238*α0807
  α1306 = 0.5693856884667226 - 0.8371528845746959*α0807
  α1307 = 10.52865228002416 - 8.957722185570706*α0807

  k1 = integrator.fsalfirst
  k2 = f(uprev+dt*α0201*k1, p, t + c2*dt)
  k3 = f(uprev+dt*(α0301*k1+α0302*k2), p, t + c3*dt)
  k4 = f(uprev+dt*(α0401*k1       +α0403*k3), p, t + c4*dt)
  k5 = f(uprev+dt*(α0501*k1       +α0503*k3+α0504*k4), p, t + c5*dt)
  k6 = f(uprev+dt*(α0601*k1                +α0604*k4+α0605*k5), p, t + c6*dt)
  k7 = f(uprev+dt*(α0701*k1                +α0704*k4+α0705*k5+α0706*k6), p, t + c7*dt)
  k8 = f(uprev+dt*(α0801*k1                +α0804*k4+α0805*k5+α0806*k6+α0807*k7), p, t + c8*dt)
  k9 = f(uprev+dt*(α0901*k1                +α0904*k4+α0905*k5+α0906*k6+α0907*k7+α0908*k8), p, t + c9*dt)
  k10 =f(uprev+dt*(α1001*k1                +α1004*k4+α1005*k5+α1006*k6+α1007*k7+α1008*k8+α1009*k9), p, t + c10*dt)
  k11= f(uprev+dt*(α1101*k1                +α1104*k4+α1105*k5+α1106*k6+α1107*k7+α1108*k8+α1109*k9+α1110*k10), p, t + c11*dt)
  k12= f(uprev+dt*(α1201*k1                +α1204*k4+α1205*k5+α1206*k6+α1207*k7+α1208*k8+α1209*k9+α1210*k10+α1211*k11), p, t+c12*dt)
  k13= f(uprev+dt*(α1301*k1                +α1304*k4+α1305*k5+α1306*k6+α1307*k7+α1308*k8+α1309*k9+α1310*k10+α1311*k11), p, t+c13*dt)
  integrator.destats.nf += 12
  u = uprev + dt*(β1*k1+β6*k6+β7*k7+β8*k8+β9*k9+β10*k10+β11*k11+β12*k12+β13*k13)
  if integrator.opts.adaptive
    utilde = dt*(β1tilde*k1 + β6tilde*k6 + β7tilde*k7 + β8tilde*k8 + β9tilde*k9 + β10tilde*k10 + β11tilde*k11 + β12tilde*k12)
    atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
  integrator.fsallast = f(u, p, t+dt)
  integrator.destats.nf += 1
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.u = u
end

function initialize!(integrator, cache::PFRK87Cache)
  integrator.fsalfirst = cache.fsalfirst
  integrator.fsallast = cache.k
  integrator.kshortsize = 2
  resize!(integrator.k, integrator.kshortsize)
  integrator.k[1] = integrator.fsalfirst
  integrator.k[2] = integrator.fsallast
  integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
  integrator.destats.nf += 1
end

@muladd function perform_step!(integrator, cache::PFRK87Cache, repeat_step=false)
  @unpack t,dt,uprev,u,f,p = integrator
  @unpack α0201, α0301, α0401, α0501, α0601, α0701, α0302, α0403, α0503, α0504, α0604, α0704, α0605, α0705, α0706, α0908, α1008, α1108, α1208, α1308, α1009, α1109, α1209, α1309, α1110, α1210, α1310, α1211, α1311, β1, β6, β7, β8, β9, β10, β11, β12, β13, β1tilde, β6tilde, β7tilde, β8tilde, β9tilde, β10tilde, β11tilde, β12tilde, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13 = cache.tab
  @unpack k2,k3,k4,k5,k6,k7,k8,k9,k10,k11,k12,k13,utilde,tmp,atmp,k = cache

  alg = unwrap_alg(integrator, true)
  ν = alg.omega*dt
  νsq = ν^2

  c_ν = -0.19781108078634084 - (0.164050909125528499*νsq) + (0.042578310088756321*(νsq^2)) - (0.002300513610963998*(νsq^3)) + (0.000033467244551879287*(νsq^4)) - (7.8661142036921924*(10^(-8))*(νsq^5))
  d_ν = 1-(0.296457092123567400*νsq) + (0.0015793885907465726*(νsq^2)) - (0.00018913011771688527*(νsq^3)) + (0.000017089234650765179*(νsq^4)) - (1.2705211682518626*(10^(-7))*(νsq^5))

  α0807 = c_ν/d_ν
  α0801 = 0.026876256 + 0.0576576*α0807
  α0804 = 0.22464336 - 0.944944*α0807
  α0805 = 0.000369024 - 0.2061696*α0807
  α0806 = 0.21311136 + 0.093456*α0807
  α0901 = 0.07239997637512857 + 0.01913119863380767*α0807
  α0904 = -0.688400520601143 - 0.3135390887207368*α0807
  α0905 = -0.17301267570583073 - 0.06840852844816077*α0807
  α0906 = 0.1440060555560846 + 0.031009360422930017*α0807
  α0907 = 0.9982362892760762 + 0.33180705811215994*α0807
  α1001 = 0.16261514523236525 - 0.12125171966747463*α0807
  α1004 = -2.1255544052061124 + 1.9871809612169453*α0807
  α1005 = -0.216403903283323 + 0.43356675517460624*α0807
  α1006 = -0.060417230254934076 - 0.1965343807796979*α0807
  α1007 = 2.4846281621788395 - 2.102961615944379*α0807
  α1101 = -1.0320124180911034 + 1.061943768952537*α0807
  α1104 = 13.666683232895137 - 17.40407843561103*α0807
  α1105 = 0.25990355211486116 - 3.797253476860588*α0807
  α1106 = -5.759316475814002 + 1.7212824826428488*α0807
  α1107 = -12.822511612651839 + 18.41810566087623*α0807
  α1201 = 0.2478349764611783 - 0.06383934946543009*α0807
  α1204 = -4.593782880309185 + 1.046256005127882*α0807
  α1205 = -0.39566692537411896 + 0.22827403748244698*α0807
  α1206 = -3.0673550479691665 - 0.10347586863902129*α0807
  α1207 = 5.386688702227177 - 1.1072148245058775*α0807
  α1301 = 0.7332242174431163 - 0.5164807626867616*α0807
  α1304 = -10.196728938160977 + 8.464545832921925*α0807
  α1305 = -0.43865244706547707 + 1.846809999910238*α0807
  α1306 = 0.5693856884667226 - 0.8371528845746959*α0807
  α1307 = 10.52865228002416 - 8.957722185570706*α0807

  k1 = cache.fsalfirst
  f(k1, uprev, p, t)
  @.. tmp = uprev+dt*α0201*k1
  f(k2, tmp, p, t + c2*dt)
  @.. tmp = uprev+dt*(α0301*k1+α0302*k2)
  f(k3, tmp, p, t + c3*dt)
  @.. tmp = uprev+dt*(α0401*k1+α0403*k3)
  f(k4, tmp, p, t + c4*dt)
  @.. tmp = uprev+dt*(α0501*k1+α0503*k3+α0504*k4)
  f(k5, tmp, p, t + c5*dt)
  @.. tmp = uprev+dt*(α0601*k1+α0604*k4+α0605*k5)
  f(k6, tmp, p, t + c6*dt)
  @.. tmp = uprev+dt*(α0701*k1+α0704*k4+α0705*k5+α0706*k6)
  f(k7, tmp, p, t + c7*dt)
  @.. tmp = uprev+dt*(α0801*k1+α0804*k4+α0805*k5+α0806*k6+α0807*k7)
  f(k8, tmp, p, t + c8*dt)
  @.. tmp = uprev+dt*(α0901*k1+α0904*k4+α0905*k5+α0906*k6+α0907*k7+α0908*k8)
  f(k9, tmp, p, t + c9*dt)
  @.. tmp = uprev+dt*(α1001*k1+α1004*k4+α1005*k5+α1006*k6+α1007*k7+α1008*k8+α1009*k9)
  f(k10, tmp, p, t + c10*dt)
  @.. tmp = uprev+dt*(α1101*k1+α1104*k4+α1105*k5+α1106*k6+α1107*k7+α1108*k8+α1109*k9+α1110*k10)
  f(k11, tmp, p, t + c11*dt)
  @.. tmp = uprev+dt*(α1201*k1+α1204*k4+α1205*k5+α1206*k6+α1207*k7+α1208*k8+α1209*k9+α1210*k10+α1211*k11)
  f(k12, tmp, p, t+c12*dt)
  @.. tmp = uprev+dt*(α1301*k1+α1304*k4+α1305*k5+α1306*k6+α1307*k7+α1308*k8+α1309*k9+α1310*k10+α1311*k11)
  f(k13, tmp, p, t+c13*dt)
  @.. u = uprev + dt*(β1*k1+β6*k6+β7*k7+β8*k8+β9*k9+β10*k10+β11*k11+β12*k12+β13*k13)
  integrator.destats.nf += 13
  if integrator.opts.adaptive
    @.. utilde = dt*(β1tilde*k1 + β6tilde*k6 + β7tilde*k7 + β8tilde*k8 + β9tilde*k9 + β10tilde*k10 + β11tilde*k11 + β12tilde*k12)
    calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm,t)
    integrator.EEst = integrator.opts.internalnorm(atmp,t)
  end
  f(k, u, p, t+dt)
  integrator.destats.nf += 1
  return nothing
end


text:
using ContinuousTimePolicyGradients
using DiffEqFlux, ComponentArrays, LinearAlgebra, JLD2, OrdinaryDiffEq
using Plots

function main(maxiters_1::Int, maxiters_2::Int, Δt_save::Float32; p_NN_0 = nothing, k_a_val = 10.0)
    # model + problem parameters
    (aₐ, aₙ, bₙ, cₙ, dₙ, aₘ, bₘ, cₘ, dₘ) = Float32.([-0.3, 19.373, -31.023, -9.717, -1.948, 40.44, -64.015, 2.922, -11.803])
    (m, I_yy, S, d, ωₐ, ζₐ) = Float32.([204.02, 247.439, 0.0409, 0.2286, 150.0, 0.7])
    (g, ρ₀, H, γₐ, Rₐ, T₀, λ) = Float32.([9.8, 1.225, 8435.0, 1.4, 286.0, 288.15, 0.0065])
    (a_max, α_max, δ_max, δ̇_max, q_max, M_max, h_max) = Float32.([100.0, deg2rad(30), deg2rad(25), 1.5, deg2rad(60), 4, 11E3])

    (k_a, k_q, k_δ, k_δ̇, k_R) = Float32.([k_a_val, 0.0, 0.01, 0.1, 1E-4])

    # dynamic model
    dim_x = 7
    function dynamics_plant(t, x, u)
        (h, V, α, q, θ, δ, δ̇) = x
        δ_c = u[1]

        ρ = ρ₀ * exp(-h / H)
        Vₛ = sqrt(γₐ * Rₐ * (T₀ - λ * h))
        M = V / Vₛ
        γ = θ - α
        Q = 0.5f0 * ρ * V^2

        C_A = aₐ
        C_N = aₙ * α^3 + bₙ * α * abs(α) + cₙ * (2.0f0 - M / 3.0f0) * α + dₙ * δ
        C_M = aₘ * α^3 + bₘ * α * abs(α) + cₘ * (-7.0f0 + 8.0f0 * M / 3.0f0) * α + dₘ * δ

        α̇ = Q * S / m / V * (C_N * cos(α) - C_A * sin(α)) + g / V * cos(γ) + q

        dx = [V * sin(γ);
            Q * S / m * (C_N * sin(α) + C_A * cos(α)) - g * sin(γ);
            α̇;
            Q * S * d / I_yy * C_M;
            q;
            δ̇;
            -ωₐ^2 * (δ - δ_c) - 2.0f0 * ζₐ * ωₐ * δ̇]
        return dx
    end

    dim_x_c = 2
    function dynamics_controller(t, x_c, y, r, p_NN, policy_NN)
        (a_z, h, V, M, α, q, γ) = y
        a_z_cmd  = r[1]
        x_int    = x_c[1]
        x_ref    = x_c[2]

        y_NN = (K_A, K_I, K_R) = policy_NN([ abs(α) / α_max; M / M_max; h / h_max], p_NN)

        # dx_ref  = [-36.6667f0 -13.8889f0; 8.0f0 0.0f0] * x_ref + [4.0f0; 0.0f0] * a_z_cmd
        # a_z_ref = [-1.0083f0 3.4722f0] * x_ref
        dx_ref = (a_z_cmd - x_ref) / 0.2f0
        a_z_ref = x_ref

        dx_c = [K_A * (a_z_cmd - a_z) + q + (a_z_cmd + g * cos(γ)) / V;
                dx_ref]
        u    = [K_I * x_int + K_R * q;
                a_z_ref]
        return dx_c, u, y_NN
    end

    function dynamics_sensor(t, x)
        (h, V, α, q, θ, δ, _) = x

        ρ = ρ₀ * exp(-h / H)
        Vₛ = sqrt(γₐ * Rₐ * (T₀ - λ * h))
        M = V / Vₛ
        Q = 0.5f0 * ρ * V^2
        γ = θ - α

        C_A = aₐ
        C_N = aₙ * α^3 + bₙ * α * abs(α) + cₙ * (2.0f0 - M / 3.0f0) * α + dₙ * δ

        a_z = Q * S / m * (C_N * cos(α) - C_A * sin(α))
        y   = [a_z;
                h;
                V;
                M;
                α;
                q;
                γ]
        return y
    end

    # cost definition
    function cost_running(t, x, y, u, r)
        q       = x[4]
        δ̇       = x[7]
        a_z     = y[1]
        δ_c     = u[1]
        a_z_ref = u[2]
        a_z_cmd = r[1]
        return k_a * ((a_z - a_z_ref) / (1.0f0 + abs(a_z_cmd)))^2 + k_δ̇ * (δ̇ / δ̇_max)^2  + k_δ * (δ_c / δ_max)^2 + k_q * (q / q_max)^2
    end

    function cost_terminal(x_f, r)
        # a_z_cmd  = r[1]
        # y        = dynamics_sensor(3.0f0, x_f)
        # a_z      = y[1]
        # return ((a_z - a_z_cmd) / (1.0f0 + a_z_cmd))^2
        return 0.0f0
    end

    function cost_regularisor(p_NN)
        return k_R * norm(p_NN)^2
        # return 0.0f0
    end

    # NN construction
    dim_NN_hidden = 10
    dim_NN_input  = 3
    dim_K = 3
    K_lb  = Float32.(0.001*ones(3))
    K_ub  = Float32[4, 0.2, 2]
    policy_NN = FastChain(
        FastDense(dim_NN_input,  dim_NN_hidden, tanh),
        FastDense(dim_NN_hidden, dim_NN_hidden, tanh),
        FastDense(dim_NN_hidden, dim_K),
        (x, p) -> (K_ub - K_lb) .* σ.(x) .+ K_lb
    )

    # scenario definition
    ensemble = [ (; x₀ = Float32[h₀; V₀; zeros(5)], r = Float32[a_z_cmd])
                     for h₀      = 5E3:1E3:8E3
                     for V₀      = 7E2:1E2:9E2
                     for a_z_cmd = -1E2:2.5E1:1E2 ]
    t_span = Float32.((0.0, 3.0))
    t_save = t_span[1]:Δt_save:t_span[2]

    scenario = (; ensemble = ensemble, t_span = t_span, t_save = t_save, dim_x = dim_x, dim_x_c = dim_x_c)

    # NN training
    (result, fwd_ensemble_sol, loss_history) = CTPG_train(dynamics_plant, dynamics_controller, dynamics_sensor, cost_running, cost_terminal, cost_regularisor, policy_NN, scenario; sense_alg = InterpolatingAdjoint(autojacvec = ReverseDiffVJP(true)), ensemble_alg = EnsembleThreads(), maxiters_1 = maxiters_1, maxiters_2 = maxiters_2, opt_2 = BFGS(initial_stepnorm = 0.0001), i_nominal = 1, p_NN_0 = p_NN_0, progress_plot = false)

    return result, policy_NN, fwd_ensemble_sol, loss_history
end

## execute optimisation and simulation
@time (result, policy_NN, fwd_ensemble_sol, loss_history) = main(1000, 1000, 0.01f0; k_a_val = 100.0)

# re-execute optimisation and simulation
# p_NN_prev = result.u
# @time (result, policy_NN, fwd_ensemble_sol, loss_history) = main(1, 1000, 0.01f0; k_a_val = 100.0, p_NN_0 = p_NN_prev)

## save results
jldsave("DS_autopilot.jld2"; result, fwd_ensemble_sol, loss_history)
# p_NN_base = result.u
# jldsave("p_NN_loss_base.jld2"; p_NN_base, loss_history)

# plot simulation results
# (result, fwd_ensemble_sol, loss_history) = load(".jld2", "result", "fwd_ensemble_sol", "loss_history")

x_names = ["\$h\$" "\$V\$" "\$\\alpha\$" "\$q\$" "\$\\theta\$" "\$\\delta\$" "\$\\dot{\\delta}\$"]
vars_x = 1:6 # [1,2,3, (1,2), (2,3)]
u_names = ["\$\\delta_{c}\$"]
vars_u = 1
y_names = ["\$a_{z}\$"]
vars_y = 1
y_NN_names = ["\$K_A\$" "\$K_I\$" "\$K_R\$"]
vars_y_NN = 1:3

(f_x, f_u, f_y, f_y_NN, f_L) = view_result([], fwd_ensemble_sol, loss_history; x_names = x_names, vars_x = vars_x, u_names = u_names, vars_u = vars_u, y_names = y_names, vars_y = vars_y, y_NN_names = y_NN_names, vars_y_NN = vars_y_NN, linealpha = 0.6)

# plot gain surfaces
# (p_NN_base) = load("p_NN_loss_base.jld2", "p_NN_base")
# (a_max, α_max, δ_max, δ̇_max, q_max, M_max, h_max) = Float32.([100.0, deg2rad(30), deg2rad(25), 1.5, deg2rad(60), 4, 11E3])

# dim_NN_hidden = 10
# dim_NN_input  = 3
# dim_K = 3
# K_lb  = Float32.(0.001*ones(3))
# K_ub  = Float32[4, 0.2, 2]
# policy_NN = FastChain(
#     FastDense(dim_NN_input,  dim_NN_hidden, tanh),
#     FastDense(dim_NN_hidden, dim_NN_hidden, tanh),
#     FastDense(dim_NN_hidden, dim_K),
#     (x, p) -> (K_ub - K_lb) .* σ.(x) .+ K_lb
# )

# h = 5000.0
# α_list = 0:1E-3:deg2rad(45)
# M_list = 0.5:0.1:3.0

# func_K_A(α, M) = policy_NN([abs(α) / α_max; M / M_max; h / h_max], p_NN_base)[1]
# func_K_I(α, M) = policy_NN([abs(α) / α_max; M / M_max; h / h_max], p_NN_base)[2]
# func_K_R(α, M) = policy_NN([abs(α) / α_max; M / M_max; h / h_max], p_NN_base)[3]

# f_K_A = plot(α_list, M_list, func_K_A, st=:surface, label = :false, zlabel = "\$K_{A}\$", xlabel = "\$\\left|\\alpha\\right|\$", ylabel = "\$M\$")
# display(f_K_A)
# savefig(f_K_A, "f_K_A.pdf")

# f_K_I = plot(α_list, M_list, func_K_I, st=:surface, label = :false, zlabel = "\$K_{I}\$", xlabel = "\$\\left|\\alpha\\right|\$", ylabel = "\$M\$")
# display(f_K_I)
# savefig(f_K_I, "f_K_I.pdf")

# f_K_R = plot(α_list, M_list, func_K_R, st=:surface, label = :false, zlabel = "\$K_{R}\$", xlabel = "\$\\left|\\alpha\\right|\$", ylabel = "\$M\$")
# display(f_K_R)
# savefig(f_K_R, "f_K_R.pdf")

ext: jl


text:
# Cast any Tuple to an TupleN{T}
@inline convert_ntuple(::Type{T},d::T) where {T} = T # For zero-dimensional arrays
@inline convert_ntuple(::Type{T},d::NTuple{N,T}) where {N,T} = d
@generated function convert_ntuple(::Type{T}, d::NTuple{N,Any}) where {N,T}
    exprs = ntuple(i -> :(convert(T, d[$i])), Val(N))
    return quote
        @_inline_meta
        $(Expr(:tuple, exprs...))
    end
end

# Base gives up on tuples for promote_eltype... (TODO can we improve Base?)
@generated function promote_tuple_eltype(::Union{T,Type{T}}) where T <: Tuple
    t = Union{}
    for i = 1:length(T.parameters)
        tmp = T.parameters[i]
        if tmp <: Vararg
            tmp = tmp.parameters[1]
        end
        t = :(promote_type($t, $tmp))
    end
    return quote
        @_inline_meta
        $t
    end
end


text:
using Documenter, Evolutionary

makedocs(
    modules = [Evolutionary],
    doctest = false,
    clean = true,
    sitename = "Evolutionary.jl",
    pages = [
        "Home" => "index.md",
        "Tutorial" => "tutorial.md",
        "Constraints" => "constraints.md",
        "Algorithms" => [
            "Genetic Algorithm" => "ga.md",
            "Differential Evolution" => "de.md",
            "Evolution Strategy" => "es.md",
            "CMA-ES" => "cmaes.md",
            "Genetic Programming" => "gp.md",
            "Multi-objective EA" => "moea.md",
        ],
        "Operations" => [
            "Selection" => "selection.md",
            "Crossover" => "crossover.md",
            "Mutation" => "mutation.md",
        ],
        "Development" => "dev.md",
    ]
)

deploydocs(repo = "github.com/wildart/Evolutionary.jl.git")


text:
module RbfInterpolationTools

include("./RbfInterpolations.jl")

end


text:
using IntervalOptimisation, TaylorModels, Test, RangeEnclosures, SumOfSquares, SDPA

using DynamicPolynomials: @polyvar

available_solvers = (NaturalEnclosure(), TaylorModelsEnclosure(),
                     BranchAndBoundEnclosure(), MooreSkelboeEnclosure())

include("univariate.jl")
include("multivariate.jl")


text:
include("../aimajulia.jl");

using Test;

using Main.aimajulia;

using Main.aimajulia.utils;

#The following reinforcement learning tests are from the aima-python doctests

north, south, west, east = (1, 0), (-1, 0), (0, -1), (0, 1);

policy = Dict([Pair((1, 1), north),
                Pair((1, 2), west),
                Pair((1, 3), west),
                Pair((1, 4), west),
                Pair((2, 1), north),
                Pair((2, 3), north),
                Pair((2, 4), nothing),
                Pair((3, 1), east),
                Pair((3, 2), east),
                Pair((3, 3), east),
                Pair((3, 4), nothing)])

passive_adp_agent = PassiveADPAgentProgram(policy, aimajulia.sequential_decision_environment);
for i in 1:75
    aimajulia.run_single_trial(passive_adp_agent, aimajulia.sequential_decision_environment);
end

@test (passive_adp_agent.U[(1, 1)] > 0.15);
println("passive_adp_agent.U[(1, 1)] (expected ~0.3): ", passive_adp_agent.U[(1, 1)]);

@test (passive_adp_agent.U[(2, 1)] > 0.15);
println("passive_adp_agent.U[(2, 1)] (expected ~0.4): ", passive_adp_agent.U[(2, 1)]);

@test (passive_adp_agent.U[(1, 2)] > 0);
println("passive_adp_agent.U[(1, 2)] (expected ~0.2): ", passive_adp_agent.U[(1, 2)]);

passive_td_agent = PassiveTDAgentProgram(policy,
                                        aimajulia.sequential_decision_environment,
                                        alpha=(function(n::Number)
                                                    return (60/(59+n));
                                                end));

for i in 1:200
    aimajulia.run_single_trial(passive_td_agent, aimajulia.sequential_decision_environment);
end

@test (passive_td_agent.U[(1, 1)] > 0.15);
println("passive_td_agent.U[(1, 1)] (expected ~0.3): ", passive_td_agent.U[(1, 1)]);

@test (passive_td_agent.U[(2, 1)] > 0.15);
println("passive_td_agent.U[(2, 1)] (expected ~0.35): ", passive_td_agent.U[(2, 1)]);

@test (passive_td_agent.U[(1, 2)] > 0.13);
println("passive_td_agent.U[(1, 2)] (expected ~0.25): ", passive_td_agent.U[(1, 2)]);

qlearning_agent = QLearningAgentProgram(aimajulia.sequential_decision_environment,
                                        5,
                                        2,
                                        alpha=(function(n::Number)
                                                    return (60/(59 + n));
                                                end));

for i in 1:200
    aimajulia.run_single_trial(qlearning_agent, aimajulia.sequential_decision_environment);
end

@test (qlearning_agent.Q[((2, 1), (1, 0))]>= -0.5);
println("qlearning_agent.Q[((2, 1), (1, 0))] expected (0.1): ", qlearning_agent.Q[((2, 1), (1, 0))]);

@test (qlearning_agent.Q[((1, 2), (-1, 0))] <= 0.5);
println("qlearning_agent.Q[((1, 2), (-1, 0))] expected (-0.1): ", qlearning_agent.Q[((1, 2), (-1, 0))]);

max_stars_repo_name: mikhail-j/aima-julia-aimacode-fork

max_issues_repo_name: mikhail-j/aima-julia-aimacode-fork

max_forks_repo_name: mikhail-j/aima-julia-aimacode-fork


text:
# Package for computing multiphase pressure profiles for gas lift optimization of oil & gas wells

module PressureDrop

using Requires

import Base.show #to export type printing methods

export  Wellbore, GasliftValves, WellModel, read_survey, read_valves,
        pressure_atmospheric,
        traverse_topdown, casing_traverse_topdown, pressure_and_temp!, pressures_and_temp!, gaslift_model!,
        plot_pressure, plot_pressures, plot_temperature, plot_pressureandtemp, plot_gaslift,
        valve_calcs, valve_table, estimate_valve_Rvalue,
        BeggsAndBrill,
        HagedornAndBrown,
        Shiu_wellboretemp, Ramey_temp, Shiu_Beggs_relaxationfactor, linear_wellboretemp,
        LeeGasViscosity,
        HankinsonWithWichertPseudoCriticalTemp,
        HankinsonWithWichertPseudoCriticalPressure,
        PapayZFactor,
        KareemEtAlZFactor,
        KareemEtAlZFactor_simplified,
        StandingSolutionGOR, StandingBubblePoint,
        StandingOilVolumeFactor,
        BeggsAndRobinsonDeadOilViscosity,
        GlasoDeadOilViscosity,
        ChewAndConnallySaturatedOilViscosity,
        GouldWaterVolumeFactor,
        SerghideFrictionFactor,
        ChenFrictionFactor

const pressure_atmospheric = 14.7 #used to adjust calculations between psia & psig

include("types.jl")
include("utilities.jl")
include("pvtproperties.jl")
include("valvecalculations.jl")
include("pressurecorrelations.jl")
include("tempcorrelations.jl")
include("casingcalculations.jl")

#lazy loading for Gadfly:
function __init__()
    @require Gadfly = "c91e804a-d5a3-530f-b6f0-dfbca275c004" include("plottingfunctions.jl")
end

#strip args from a struct and pass as kwargs to a function:
macro run(input, func)

    return quote
        local var = $(esc(input)) #resolve using the macro call environment
        local fn = $(esc(func))

        fields = fieldnames(typeof(var))
        values = map(f -> getfield(var, f), fields)

        args = (;(f=>v for (f,v) in zip(fields,values))...) #splat and append to convert to NamedTuple that can be passed as kwargs

        fn(;args...)

    end
end

#%% core functions
"""
`calculate_pressuresegment(<arguments>)`

Pressure inputs are in **psia**.

Helper function to calculate the pressure drop for a single pipe segment, using an outlet-referenced approach.

Method (fixed-point iteration to account for pressure dependence of fluid PVT properties) :
1. Generates PVT properties for the average conditions in the segment for a given estimate for the pressure drop in the second.
2. Calculates a new estimated pressure drop using the PVT properties.
3. Compares the original and new pressure drop estimates to validate the stability of the estimate.
4. Iterate through steps 1-3 until a stable estimate is found (indicated by obtaining a difference between pre- and post-PVT that is within the given error tolerance).

"""
function calculate_pressuresegment(pressurecorrelation::Function, p_initial, dp_est, t_avg,
                                            dh_md, dh_tvd, inclination, uphill_flow, id, roughness,
                                            q_o, q_w, GLR, R_b, APIoil, sg_water, sg_gas,
                                            Z_correlation::Function, P_pc, T_pc,
                                            gas_viscosity_correlation::Function, solutionGORcorrelation::Function, bubblepoint, oilVolumeFactor_correlation::Function, waterVolumeFactor_correlation::Function,
                                            dead_oil_viscosity_correlation::Function, live_oil_viscosity_correlation::Function, frictionfactor::Function, error_tolerance = 0.1)

    function dP(dp) #calculate delta pressure dP as a function of input pressure estimate dp
        p_avg = p_initial + dp/2

        Z = Z_correlation(P_pc, T_pc, p_avg, t_avg)
        ρ_g = gasDensity_insitu(sg_gas, Z, p_avg, t_avg)
        B_g = gasVolumeFactor(p_avg, Z, t_avg)
        μ_g = gas_viscosity_correlation(sg_gas, p_avg, t_avg, Z)
        R_s = solutionGORcorrelation(APIoil, sg_gas, p_avg, t_avg, R_b, bubblepoint)
        v_sg = gasvelocity_superficial(q_o, q_w, GLR, R_s, id, B_g)
        B_o = oilVolumeFactor_correlation(APIoil, sg_gas, R_s, p_avg, t_avg)
        B_w = waterVolumeFactor_correlation(p_avg, t_avg)
        v_sl = liquidvelocity_superficial(q_o, q_w, id, B_o, B_w)
        ρ_l = mixture_properties_simple(q_o, q_w, oilDensity_insitu(APIoil,  sg_gas,  R_s,  B_o), waterDensity_insitu(sg_water, B_w))
        σ_l = mixture_properties_simple(q_o, q_w, gas_oil_interfacialtension(APIoil, p_avg, t_avg), gas_water_interfacialtension(p_avg, t_avg))
        μ_oD = dead_oil_viscosity_correlation(APIoil, t_avg)
        μ_l = mixture_properties_simple(q_o, q_w, live_oil_viscosity_correlation(μ_oD, R_s), assumedWaterViscosity)

        return pressurecorrelation(dh_md, dh_tvd, inclination, id,
                                        v_sl, v_sg, ρ_l, ρ_g, σ_l, μ_l, μ_g, roughness, p_avg, frictionfactor,
                                        uphill_flow)
    end

    dp_calc = dP(dp_est)
    while abs(dp_est - dp_calc) > error_tolerance

        dp_est = dp_calc
        dp_calc = dP(dp_est)
    end

    return dp_calc #allows negatives
end

"""
`traverse_topdown(;<named arguments>)`

Develop pressure traverse from wellhead down to datum in psia, returning a pressure profile as an Array{Float64,1}.

Pressure correlation functions available:
- `BeggsAndBrill` with Payne correction factors
- `HagedornAndBrown` with Griffith and Wallis bubble flow correction

# Arguments

All arguments are named keyword arguments.

## Required
- `wellbore::Wellbore`: Wellbore object that defines segmentation/mesh, with md, tvd, inclination, and hydraulic diameter
- `roughness`: pipe wall roughness in inches
- `temperatureprofile::Array{Float64, 1}`: temperature profile (in °F) as an array with **matching entries for each pipe segment defined in the Wellbore input**
- `WHP`: outlet pressure (wellhead pressure) in **psig**
- `dp_est`: estimated starting pressure differential (in psi) to use for all segments--impacts convergence time
- `q_o`: oil rate in stocktank barrels/day
- `q_w`: water rate in stb/d
- `GLR`: **total** wellhead gas:liquid ratio, inclusive of injection gas, in scf/bbl
- `APIoil`: API gravity of the produced oil
- `sg_water`: specific gravity of produced water
- `sg_gas`: specific gravity of produced gas

## Optional
- `injection_point = missing`: injection point in MD for gas lift, above which total GLR is used, and below which natural GLR is used
- `naturalGLR = missing`: GLR to use below point of injection, in scf/bbl
- `pressurecorrelation::Function = BeggsAndBrill: pressure correlation to use
- `error_tolerance = 0.1`: error tolerance for each segment in psi
- `molFracCO2 = 0.0`, `molFracH2S = 0.0`: produced gas fractions of hydrogen sulfide and CO2, [0,1]
- `pseudocrit_pressure_correlation::Function = HankinsonWithWichertPseudoCriticalPressure`: psuedocritical pressure function to use
- `pseudocrit_temp_correlation::Function = HankinsonWithWichertPseudoCriticalTemp`: pseudocritical temperature function to use
- `Z_correlation::Function = KareemEtAlZFactor`: natural gas compressibility/Z-factor correlation to use
- `gas_viscosity_correlation::Function = LeeGasViscosity`: gas viscosity correlation to use
- `solutionGORcorrelation::Function = StandingSolutionGOR`: solution GOR correlation to use
- `bubblepoint::Union{Function, Real} = StandingBubblePoint`: either bubble point correlation or bubble point in **psia**
- `oilVolumeFactor_correlation::Function = StandingOilVolumeFactor`: oil volume factor correlation to use
- `waterVolumeFactor_correlation::Function = GouldWaterVolumeFactor`: water volume factor correlation to use
- `dead_oil_viscosity_correlation::Function = GlasoDeadOilViscosity`: dead oil viscosity correlation to use
- `live_oil_viscosity_correlation::Function = ChewAndConnallySaturatedOilViscosity`: saturated oil viscosity correction function to use
- `frictionfactor::Function = SerghideFrictionFactor`: correlation function for Darcy-Weisbach friction factor
"""
function traverse_topdown(;wellbore::Wellbore, roughness, temperatureprofile::Array{Float64, 1},
                            pressurecorrelation::Function = BeggsAndBrill,
                            WHP, dp_est, error_tolerance = 0.1,
                            q_o, q_w, GLR, injection_point = missing, naturalGLR = missing,
                            APIoil, sg_water, sg_gas, molFracCO2 = 0.0, molFracH2S = 0.0,
                            pseudocrit_pressure_correlation::Function = HankinsonWithWichertPseudoCriticalPressure, pseudocrit_temp_correlation::Function = HankinsonWithWichertPseudoCriticalTemp,
                            Z_correlation::Function = KareemEtAlZFactor, gas_viscosity_correlation::Function = LeeGasViscosity, solutionGORcorrelation::Function = StandingSolutionGOR, bubblepoint = StandingBubblePoint,
                            oilVolumeFactor_correlation::Function = StandingOilVolumeFactor, waterVolumeFactor_correlation::Function = GouldWaterVolumeFactor,
                            dead_oil_viscosity_correlation::Function = GlasoDeadOilViscosity, live_oil_viscosity_correlation::Function = ChewAndConnallySaturatedOilViscosity, frictionfactor::Function = SerghideFrictionFactor,
                            kwargs...) #catch extra arguments from a WellModel for convenience

    @assert q_o >= 0. && q_w >= 0. && GLR >= 0. && (naturalGLR === missing || (GLR >= naturalGLR >= 0.)) "Negative rates or NGLR > GLR not supported."

    WHP += pressure_atmospheric #convert psig input to psia for internal PVT functions

    nsegments = length(wellbore.md)

    @assert nsegments == length(temperatureprofile) "Number of wellbore segments does not match number of temperature points."

    if !ismissing(injection_point) && !ismissing(naturalGLR)
        inj_index = searchsortedlast(wellbore.md, injection_point)

        if wellbore.md[inj_index] != injection_point
            if (injection_point - wellbore.md[inj_index]) > (wellbore.md[inj_index+1] - injection_point) #choose closest point
                inj_index += 1
            end

            @info """Specified injection point at $injection_point' MD not explicitly included in wellbore. Using $(round(wellbore.md[inj_index],digits=1))' MD as an approximate match.
            Use the Wellbore constructor with a set of gas lift valves to add precise injection points."""
        end

        GLRs = vcat(repeat([GLR], inner = inj_index), repeat([naturalGLR], inner = nsegments - inj_index))
        R_b = repeat([naturalGLR], inner = nsegments) #reservoir total solution GOR above bpp
    elseif !ismissing(injection_point) || !ismissing(naturalGLR)
        @info "Both an injection point and natural GLR should be specified--ignoring partial specification."
        GLRs = repeat([GLR], inner = nsegments)
        R_b = GLRs
    else #no injection point
        GLRs = repeat([GLR], inner = nsegments)
        R_b = GLRs
    end

    pressures = Array{Float64, 1}(undef, nsegments)
    pressure_initial = pressures[1] = WHP

    P_pc = pseudocrit_pressure_correlation(sg_gas, molFracCO2, molFracH2S)
    _, T_pc, _ = pseudocrit_temp_correlation(sg_gas, molFracCO2, molFracH2S)

    @inbounds for i in 2:nsegments
        inclination = (wellbore.inc[i] + wellbore.inc[i-1])/2

        dp_est = calculate_pressuresegment(pressurecorrelation, pressure_initial, dp_est,
                                                    (temperatureprofile[i] + temperatureprofile[i-1])/2, #average temperature
                                                    wellbore.md[i] - wellbore.md[i-1], #dh_md
                                                    wellbore.tvd[i] - wellbore.tvd[i-1], #dh_tvd
                                                    inclination, #average inclination between survey points
                                                    inclination <= 90.0, #uphill_flow
                                                    wellbore.id[i], roughness,
                                                    q_o, q_w, GLRs[i], R_b[i], APIoil, sg_water, sg_gas,
                                                    Z_correlation, P_pc, T_pc,
                                                    gas_viscosity_correlation, solutionGORcorrelation, bubblepoint, oilVolumeFactor_correlation, waterVolumeFactor_correlation,
                                                    dead_oil_viscosity_correlation, live_oil_viscosity_correlation, frictionfactor, error_tolerance)

        pressure_initial += dp_est
        pressures[i] = pressure_initial
    end

    return pressures .- pressure_atmospheric #convert back to psig for user-facing output
end

"""
`traverse_topdown(;model::WellModel)`

calculate top-down traverse from a WellModel object. Requires the following fields to be defined in the model:

...
"""
function traverse_topdown(model::WellModel)

    @run model traverse_topdown
end

"""
BHP_summary(pressures, well)

Print the summary for a bottomhole pressure traverse of a well.
"""
function BHP_summary(pressures, well)
    println("Flowing bottomhole pressure of $(round(pressures[end], digits = 1)) psig at $(well.md[end])' MD.",
        "\nAverage gradient $(round(pressures[end]/well.md[end], digits = 3)) psi/ft (MD), $(round(pressures[end]/well.tvd[end], digits = 3)) psi/ft (TVD).")
end

"""
`pressure_and_temp(;model::WellModel)`

Develop pressure traverse in psia and temperature profile in °F from wellhead down to datum for a WellModel object. Requires the following fields to be defined in the model:

Returns a pressure profile as an Array{Float64,1} and updates the passed WellModel's temperature profile, referenced to the measured depths in the original Wellbore object.

# Arguments

All arguments are defined in the model object; see the `WellModel` documentation for reference.

Pressure correlation functions available:
- `BeggsAndBrill` with Payne correction factors
- `HagedornAndBrown` with Griffith and Wallis bubble flow correction

Temperature methods available:
- "Shiu" to utilize the Ramey 1962 method with the Shiu 1980 relaxation factor correlation
- "linear" for a linear interpolation between wellhead and bottomhole temperature based on TVD

## Required `WellModel` fields
- `well::Wellbore`: Wellbore object that defines segmentation/mesh, with md, tvd, inclination, and hydraulic diameter
- `roughness`: pipe wall roughness in inches
- `temperature_method = "linear"`: temperature method to use; "Shiu" for Ramey method with Shiu relaxation factor, "linear" for linear interpolation
- `WHT = missing`: wellhead temperature in °F; required for `temperature_method = "linear"`
- `geothermal_gradient = missing`: geothermal gradient in °F per 100 ft; required for `temperature_method = "Shiu"`
- `BHT` = bottomhole temperature in °F
- `WHP`: absolute outlet pressure (wellhead pressure) in **psig**
- `dp_est`: estimated starting pressure differential (in psi) to use for all segments--impacts convergence time
- `q_o`: oil rate in stocktank barrels/day
- `q_w`: water rate in stb/d
- `GLR`: **total** wellhead gas:liquid ratio, inclusive of injection gas, in scf/bbl
- `APIoil`: API gravity of the produced oil
- `sg_water`: specific gravity of produced water
- `sg_gas`: specific gravity of produced gas

## Optional `WellModel` fields
- `injection_point = missing`: injection point in MD for gas lift, above which total GLR is used, and below which natural GLR is used
- `naturalGLR = missing`: GLR to use below point of injection, in scf/bbl
- `pressurecorrelation::Function = BeggsAndBrill: pressure correlation to use
- `error_tolerance = 0.1`: error tolerance for each segment in psi
- `molFracCO2 = 0.0`, `molFracH2S = 0.0`: produced gas fractions of hydrogen sulfide and CO2, [0,1]
- `pseudocrit_pressure_correlation::Function = HankinsonWithWichertPseudoCriticalPressure`: psuedocritical pressure function to use
- `pseudocrit_temp_correlation::Function = HankinsonWithWichertPseudoCriticalTemp`: pseudocritical temperature function to use
- `Z_correlation::Function = KareemEtAlZFactor`: natural gas compressibility/Z-factor correlation to use
- `gas_viscosity_correlation::Function = LeeGasViscosity`: gas viscosity correlation to use
- `solutionGORcorrelation::Function = StandingSolutionGOR`: solution GOR correlation to use
- `bubblepoint::Union{Function, Real} = StandingBubblePoint`: either bubble point correlation or bubble point in **psia**
- `oilVolumeFactor_correlation::Function = StandingOilVolumeFactor`: oil volume factor correlation to use
- `waterVolumeFactor_correlation::Function = GouldWaterVolumeFactor`: water volume factor correlation to use
- `dead_oil_viscosity_correlation::Function = GlasoDeadOilViscosity`: dead oil viscosity correlation to use
- `live_oil_viscosity_correlation::Function = ChewAndConnallySaturatedOilViscosity`: saturated oil viscosity correction function to use
- `frictionfactor::Function = SerghideFrictionFactor`: correlation function for Darcy-Weisbach friction factor
- `outlet_referenced = true`: whether to use outlet pressure (WHP) or inlet pressure (BHP) for starting point
"""
function pressure_and_temp!(m::WellModel, summary = true)

    if m.temperature_method == "linear"
        @assert !(any(ismissing.((m.WHT, m.BHT)))) "Must specific a wellhead temperature & BHT to utilize linear temperature method."
        m.temperatureprofile = @run m linear_wellboretemp
    elseif m.temperature_method == "Shiu"
        @assert !(any(ismissing.((m.BHT, m.geothermal_gradient)))) "Must specify a geothermal gradient & BHT to utilize Shiu/Ramey temperature method.\nRefer to published geothermal gradient maps for your region to establish a sensible default."
        m.temperatureprofile = @run m Shiu_wellboretemp
    else
        throw(ArgumentError("Invalid temperature method. Use one of (\"Shiu\", \"linear\")."))
    end

    pressures = traverse_topdown(m)
    summary ? BHP_summary(pressures, m.wellbore) : nothing

    return pressures
end

"""
`pressures_and_temp!(m::WellModel)`

Returns a tubing pressure profile as an Array{Float64,1}, casing pressure profile as an Array{Float64,1}, and updates the passed WellModel's temperature profile, referenced to the measured depths in the original Wellbore object.

# Arguments

See `WellModel` documentation.
"""
function pressures_and_temp!(m::WellModel, summary = true)

    tubing_pressures = pressure_and_temp!(m, summary)
    casing_pressures = casing_traverse_topdown(m)

    return tubing_pressures, casing_pressures
end

"""
`gaslift_model!(m::WellModel; find_injectionpoint::Bool = false, dp_min = 100)`

Returns a tubing pressure profile as an Array{Float64,1}, casing pressure profile as an Array{Float64,1}, valve data table, and updates the passed WellModel's temperature profile,

# Arguments

See `WellModel` documentation.

- `find_injectionpoint::Bool = false`: whether to automatically infer the injection point (taken as the lowest reasonable point of lift based on differential pressure)*
- `dp_min = 100`: minimum casing-tubing differential pressure at depth to infer an injection point

*"greedy opening" heuristic: select _lowest_ non-orifice valve where CP @ depth is within operating envelope (below opening pressure but still above closing pressure) and has greater than the indicated differential pressure (`dp_min`)
"""
function gaslift_model!(m::WellModel, summary = false; find_injectionpoint::Bool = false, dp_min = 100)

    if find_injectionpoint
        m.injection_point = m.wellbore.md[end]
    elseif m.injection_point === missing || m.naturalGLR === missing
        @info "Performing gas lift calculations without defined injection information (point of injection, natural GLR) and without falling back to a calculated injection point."
    end

    tubing_pressures, casing_pressures = pressures_and_temp!(m, false);
    valvedata, injection_depth = valve_calcs(valves = m.valves, well = m.wellbore, sg_gas = m.sg_gas_inj, tubing_pressures = tubing_pressures, casing_pressures = casing_pressures, tubing_temps = m.temperatureprofile, casing_temps = m.temperatureprofile .* m.casing_temp_factor,
                            dp_min = dp_min)

    #currently doesn't account for changing temp profile
    if find_injectionpoint
        @info "Inferred injection depth @ $injection_depth' MD."
        m.injection_point = injection_depth
        tubing_pressures = traverse_topdown(m)
        casing_pressures = casing_traverse_topdown(m)
        valvedata, _ = valve_calcs(valves = m.valves, well = m.wellbore, sg_gas = m.sg_gas_inj, tubing_pressures = tubing_pressures, casing_pressures = casing_pressures, tubing_temps = m.temperatureprofile, casing_temps = m.temperatureprofile .* m.casing_temp_factor,
                                dp_min = dp_min)
    end

    summary ? BHP_summary(tubing_pressures, m.wellbore) : nothing
    return tubing_pressures, casing_pressures, valvedata
end

end #module PressureDrop


text:
using Altro
using TrajectoryOptimization
using StaticArrays
using BenchmarkTools
using RobotZoo
using RobotDynamics
using StaticArrays, LinearAlgebra
using ForwardDiff
using Test
using FiniteDiff
using Random
const TO = TrajectoryOptimization
const RD = RobotDynamics

# Test Functions
function Πsoc(x)
    v = x[1:end-1]
    s = x[end]
    a = norm(v)
    if a <= -s
        return zero(x)
    elseif a <= s
        return x
    elseif a >= abs(s)
        x̄ = append!(v, a)
        return 0.5*(1 + s/a) * x̄
    end
    throw(ErrorException("Invalid second-order cone"))
end
function in_soc(x)
    v = x[1:end-1]
    s = x[end]
    a = norm(v)
	a <= s
end

function alcost(x,z,μ)
    c = push(x, u_bnd)
    proj = Πsoc(z - μ .* c)
    1/2μ*(proj'proj - z'z)
end

function auglag(f,g, x,z,μ)
    # ceq = h(x)
    L0 = f(x) #+ y'ceq + 0.5*μ*ceq'ceq

    # cones = g(x)
    # p = length(cones)
    for i = 1:length(z)
        L0 += alcost(x[uinds[i]], z[i], μ)
        # cone = cones[i]
		# proj = Πsoc(z[i] - μ*cones[i])
        # Iμ = I * inv(μ)
		# pen = proj'proj - z[i]'z[i]
        # L0 += 1 / 2μ * pen
    end
    return L0
end

function DI_objective(Q,R,Qf,xf,dt)
    g = Float64[]
    h = Float64[]
    c = 0.0
    for k = 1:N-1
        append!(g, -Q*xf)
        append!(g, zeros(m))
        append!(h, Q.diag)
        append!(h, R.diag)
        c += 0.5*xf'Q*xf
	end
	# g .*= dt
	# h .*= dtk
	# c *= dt
    append!(g, -Qf*xf)
    # append!(g, zeros(m))
    append!(h, Qf.diag)
    # append!(h, R.diag)
    c += 0.5*xf'Qf*xf
    H = Diagonal(h)
    return H,g,c
end

function DI_dynamics(x0, N; dt=0.1, gravity::Bool=false)
    n = length(x0)
    m = n ÷ 2
    D = m
    NN = N*(n+m)

    #  Continuous dynamics
    Ac = zeros(n,n)
    Bc = zeros(n,m)
    for i = 1:D
        Ac[i,D+i] = 1
        Bc[D+i,i] = 1
    end

    # Euler integration
    Ad = I + Ac*dt
    Bd = Bc*dt

    iA = 1:n
    jA = iA
    jB = n .+ (1:m)
    jA2 = jA .+ (n+m)

    A = zeros(N*n, NN)
    b = zeros(N*n)
    A[iA,jA] .= I(n)
    b[iA] .= -x0
    iA = iA .+ n
    for k = 1:N-1
        A[iA,jA] .= Ad
        A[iA,jB] .= Bd
        A[iA,jA2] .= -I(n)
        if gravity
            b[iA[D]] = -9.81
        end

        # Advance inds
        iA = iA .+ n
        jA = jA .+ (n+m)
        jB = jB .+ (n+m)
        jA2 = jA2 .+ (n+m)
    end
    return A,b
end

function DI_cones(::Val{D},N) where D
    uinds = SVector{D}(1:D) .+ 2D
    qinds = [uinds .+ (k-1)*3D for k = 1:N-1]
end
DI_cones(D,N) = DI_cones(Val(D),N)  # WARNING: type unstable

## Setup
D,N = 2,11
n,m = 2D,D
dt = 0.1
tf = (N-1)*dt
x0 = @SVector fill(0., n)
xf = [(@SVector fill(1.,D)); (@SVector fill(0, D))]
Q = Diagonal(@SVector fill(1.0, n)) * dt
R = Diagonal(@SVector fill(1e-1, m)) * dt
Qf = (N-1)*Q * 100
qinds = DI_cones(D,N)
u_bnd = 6.0

# Define batch functions
H_obj,g_obj,c_obj = DI_objective(Q,R,Qf,xf,dt)
Adyn,bdyn = DI_dynamics(x0,N, dt=dt)

di_obj(x) = 0.5*x'H_obj*x + g_obj'x + c_obj
di_dyn(x) = Adyn*x + bdyn
di_soc(x) = [push(x[qi],u_bnd) for qi in qinds]

# Define using ALTRO
model = RobotZoo.DoubleIntegrator(D)

obj = LQRObjective(Q, R, Qf, xf, N)
cons = ConstraintList(n,m,N)
cone = TO.SecondOrderCone()
# cone = TO.Inequality()
u_bnd = 6.0
connorm = NormConstraint(n, m, u_bnd, cone, :control)
add_constraint!(cons, connorm, 1:N-1)
prob = Problem(model, obj, x0, tf, xf=xf, constraints=cons, integration=RD.Euler(model))
solver = Altro.ALSolver(prob, use_static=Val(true))
conSet = TO.get_constraints(solver)
alobj = TO.get_objective(solver)

# Initialization
inds = LinearIndices(zeros(3D,N))
xinds = [SVector{n}(z[1:n]) for z in eachcol(inds)]
uinds = [SVector{m}(z[n+1:end]) for z in eachcol(inds)][1:end-1]
NN = 3D*N-D
P = (N-1)*m
x0 = rand(NN)
z = rand.(length.(di_soc(x0)))
z0 = deepcopy(z)
μ = 2.4

X0 = [x0[xi] for xi in xinds]
U0 = [x0[ui] for ui in uinds]
Z0 = SampledTrajectory(X0,U0, dt=dt)
initial_trajectory!(solver, Z0)

## Test objective and soc constraints
z = deepcopy(z0)
Altro.reset_duals!(conSet)
@test cost(alobj, Z0) ≈ di_obj(x0)
@test conSet[1].vals ≈ di_soc(x0)

# copy multipliers to ALconSet
cval = conSet[1]
for i = 1:N-1
	cval.λ[i] .= z[i]
	cval.μ[i] .= μ
end

# test augmented Lagrangian value
LA(x) = auglag(di_obj, di_soc, x, z, μ)
@test di_obj(x0) ≈ cost(alobj.obj, Z0)  # unconstrained cost
@test cost(solver) ≈ LA(x0)
@test !(cost(solver) ≈ di_obj(x0))   # make sure the AL cost isn't the same as the normal cost

λbar = z .- μ .* di_soc(x0)
statuses = [TO.cone_status(TO.SecondOrderCone(), λ) for λ in λbar]
@test :below ∈ statuses

# E = TO.QuadraticObjective(n,m,N)
E0 = Altro.CostExpansion{Float64}(n,m,N)
E = Altro.get_ilqr(solver).Efull
let solver = solver.ilqr
    Altro.cost_expansion!(solver.obj, E, solver.Z)
    @test_throws AssertionError Altro.cost_expansion!(solver.obj, E0, solver.Z)
end
# TO.cost_expansion!(E, alobj, Z0)
grad = vcat([e.grad for e in E]...)[1:end-m]
@test grad ≈ FiniteDiff.finite_difference_gradient(LA, x0)

hess_blocks = vcat([[e.xx, e.uu] for e in E]...)
hess = cat(hess_blocks..., dims=(1,2))[1:end-m, 1:end-m]
@test hess ≈ FiniteDiff.finite_difference_hessian(LA, x0) atol=1e-4

@test all(in_soc.(di_soc(x0)))
@test !all(in_soc.(z .- μ*di_soc(x0)))

# Make it go outside the cone
Random.seed!(1)
z = rand.(length.(di_soc(x0)))
z .*= 100
λbar = z .- μ .* di_soc(x0)
statuses = [TO.cone_status(TO.SecondOrderCone(), λ) for λ in λbar]
for i = 1:20
    if :outside in statuses && :in in statuses
        break
    else
        z .= rand.(length.(di_soc(x0)))
        z .*= 100
        λbar .= z .- μ .* di_soc(x0)
        statuses .= [TO.cone_status(TO.SecondOrderCone(), λ) for λ in λbar]
    end
end
@test :outside ∈ statuses
@test :in ∈ statuses
for i = 1:N-1
	cval.λ[i] .= z[i]
	cval.μ[i] .= μ
    cval.μinv[i] .= inv.(μ)
end
LA(x) = auglag(di_obj, di_soc, x, z, μ)

# E = TO.QuadraticObjective(n,m,N)
# E = Altro.CostExpansion{Float64}(n,m,N)
E = Altro.get_ilqr(solver).Efull
@test LA(x0) ≈ cost(solver, Z0)
Altro.cost_expansion!(alobj, E, Z0)
# TO.cost_expansion!(E, alobj, Z0)
alcon = conSet[1]
grad = vcat([e.grad for e in E]...)[1:end-m]
@test grad ≈ FiniteDiff.finite_difference_gradient(LA, x0) atol=1e-6
[grad FiniteDiff.finite_difference_gradient(LA, x0)]

hess_blocks = vcat([[e.xx, e.uu] for e in E]...)
hess = cat(hess_blocks..., dims=(1,2))[1:end-m, 1:end-m]
@test hess ≈ FiniteDiff.finite_difference_hessian(LA, x0) atol=1e-2

## Solve it
prob = Problem(model, obj, zero(SVector{n}), tf, xf=xf, constraints=cons, integration=RD.Euler(model))
solver = Altro.ALSolver(prob, verbose=0, show_summary=false,
	projected_newton=true, penalty_initial=100.0, penalty_scaling=50,
	cost_tolerance_intermediate=1e-1)
initial_controls!(solver, [rand(D) for k = 1:N])
solve!(solver)
@test all(abs.(norm.(controls(solver))[[1,2,N-2,N-1]] .- u_bnd) .< 1e-5)

# prob = Problem(model, obj, xf, tf, x0=x0, constraints=cons)
# solver = ALTROSolver(prob, show_summary=true, verbose=2, projected_newton=false)
# @test benchmark_solve!(solver).allocs == 0
# benchmark_solve!(solver)


text:
include("dwt.jl")
include("audiodwt.jl")
include("audiodct.jl")
using WAV
using LinearAlgebra
using Plots

# Read song, Am I Blue
# Sung by Kevin Conroy, listen for free at http://www.realkevinconroy.com/
filename = "Am I Blue.wav"
channels, Fs, nbits, opt = wavread(filename)

# Generate wavelet filters
ps = 2:6
ls = []
hs = []
ils = []
ihs = []
for p = ps
    l, h, il, ih = filtergen(p)
    push!(ls, l)
    push!(hs, h)
    push!(ils, il)
    push!(ihs, ih)
end

# Set thresholds
thresholds = 10 .^(-4:.2:-1)
wsizes = zeros(length(thresholds), length(ps))
fsizes = zeros(length(thresholds))
werrors = zeros(length(thresholds), length(ps))
ferrors = zeros(length(thresholds))

for i = 1:length(thresholds)
    threshold = thresholds[i]
    for j = 1:length(ps)
        # Fetch wavelet filters
        l = ls[j]
        h = hs[j]
        il = ils[j]
        ih = ihs[j]
        p = ps[j]

        # Compress DWT
        wcompression, wsizes[i,j] = compressaudio(channels, p, l, h, threshold)

        # Decompress DWT
        wrecovered = decompressaudio(wcompression, p, il, ih)
        werrors[i,j] = norm(wrecovered[:] - channels[:]) / norm(channels[:])
    end

    # Compress DCT
    fcompression, fsizes[i] = compressaudiodct(channels, threshold)

    # Decompress DCT
    frecovered = decompressaudiodct(fcompression)
    ferrors[i] = norm(frecovered[:] - channels[:]) / norm(channels[:])
end

# Plot
x = 100 / sizeof(channels)
plot(werrors[:,1], wsizes[:,1] .* x, label=string("DAUB", 2ps[1]));
for i = 2:length(ps)
    p = ps[i]
    plot!(werrors[:,i], wsizes[:,i] .* x, label=string("DAUB", 2p));
end
plot!(ferrors, fsizes .* x, label="DCT");
title!("Comparison of Compression Accuracy of DWT and DCT");
xaxis!("Relative Root Mean Squared Error of Reconstruction", :log10);
yaxis!("Percentage of original file size (%)")


text:
@testset "ESPIRiT" begin
  #image
  N=128
  img = shepp_logan(N)
  msk = zeros(N,N)
  msk[findall(x->x!=0,img)] .= 1

  # coil sensitivites
  smaps = birdcageSensitivity(N, 8, 1.5)
  snorm = sqrt.(sum(abs.(smaps).^2,dims=4))
  for i=1:8
    smaps[:,:,1,i] .= msk .* smaps[:,:,1,i] ./ snorm[:,:,1,1]
  end

  # simulation
  params = Dict{Symbol, Any}()
  params[:simulation] = "fast"
  params[:trajName] = "Cartesian"
  params[:numProfiles] = floor(Int64, N)
  params[:numSamplingPerProfile] = N
  params[:senseMaps] = smaps

  acqData = simulation(img, params)
  acqData = MRIReco.sample_kspace(acqData, 2.0, "poisson", calsize=15)

  ksize=(6,6) # kernel size
  ncalib = 15 # number of calibration lines
  eigThresh_1 = 0.02  # threshold for picking singular vectors of calibration matrix
  eigThresh_2 = 0.95  # threshold for eigen vector decomposition in image space

  smaps2 = espirit(acqData,ksize,ncalib,eigThresh_1=eigThresh_1,eigThresh_2=eigThresh_2)

  # evaluate error only on the supprt of smaps
  for i=1:8
    smaps2[:,:,1,i] = msk .* smaps2[:,:,1,i]
  end
  # allow for a constant phase offset
  phs = mean(angle.(smaps))
  phs2 = mean(angle.(smaps2))
  smaps2 = exp(1im*(phs-phs2)) .* smaps2

  err = norm(vec(smaps2)-vec(smaps))/norm(vec(smaps))
  @test err < 3.e-2
end


text:
using PipeLayout.JuncLoc

using JuMP
using SCIP

# don't solve nonconvex NLP to optimality
_scip = JuMP.optimizer_with_attributes(
    SCIP.Optimizer,
    "display/verblevel" => 0,
    "limits/gap" => 1e-3
)

@testset "solve junction location for three terminals (NLP)" begin
    # equilateral triangle
    nodes = [Node(0,0), Node(40,0), Node(20, sqrt(3)/2*40)]
    demand = [10, 10, -20]
    bounds = fill(Bounds(40, 80), 3)
    diams = [Diameter(t...) for t in [(0.6, 0.5), (0.8, 0.8),(1.0, 1.2)]]

    # single Steiner node with dummy location
    topo = Topology(vcat(nodes, [Node(20, 20)]),
                    [Arc(3,4), Arc(4,1), Arc(4,2)])

    solver = JuncLoc.NLP(_scip)

    @testset "little flow, smallest diameter, Steiner node in center" begin
        inst = Instance(nodes, demand, bounds, diams, ploss_coeff_nice)
        model, x, y, L, l, π = JuncLoc.make_nlp(inst, topo, solver)
        JuMP.optimize!(model)
	    status = JuMP.termination_status(model)

        @test status == MOI.OPTIMAL # includes gap limit

        xsol, ysol = JuMP.value.(x), JuMP.value.(y)
        for i=1:3 # fixed terminals
            @test xsol[i] ≈ nodes[i].x atol=0.001
            @test ysol[i] ≈ nodes[i].y atol=0.001
        end
        @test xsol[4] ≈ 20 atol=0.01
        @test ysol[4] ≈ sqrt(3)/6*40 atol=0.01

        Lsol = JuMP.value.(L)
        for i=1:3
            @test Lsol[i] ≈ sqrt(3)/3*40 atol=0.01
        end

        lsol = JuMP.value.(l)
        # smallest diameter
        @test sum(lsol[:,1]) ≈ 3 atol=0.01
        # no other
        @test sum(lsol[:, 2:end]) ≈ 0 atol=0.01
    end

    @testset "more flow, mixed diameter, Steiner node towards source" begin
        inst = Instance(nodes, 20*demand, bounds, diams, ploss_coeff_nice)
        model, x, y, L, l, π = JuncLoc.make_nlp(inst, topo, solver)
        JuMP.optimize!(model)
	    status = JuMP.termination_status(model)

        @test status == MOI.OPTIMAL

        xsol, ysol = JuMP.value.(x), JuMP.value.(y)
        for i=1:3 # fixed terminals
            @test xsol[i] ≈ nodes[i].x atol=0.001
            @test ysol[i] ≈ nodes[i].y atol=0.001
        end
        @test xsol[4] ≈ 20 atol=0.2
        @test ysol[4] >= sqrt(3)/6*40 # move near source

        Lsol = JuMP.value.(L)
        @test Lsol[1] <= sqrt(3)/3*40
        @test Lsol[2] >= sqrt(3)/3*40
        @test Lsol[3] >= sqrt(3)/3*40

        lsol = JuMP.value.(l)
        D = [d.value for d in diams]
        equiv = (lsol * D.^(-5)).^(-1/5)

        # symmetry
        @test equiv[2] ≈ equiv[3] atol=0.1
        # more flow -> larger diam?
        @test equiv[1] >= equiv[2]
    end

    @testset "using the optimize function to solve" begin
        inst = Instance(nodes, 20*demand, bounds, diams, ploss_coeff_nice)
        result = optimize(inst, topo, solver)
        @test result.status == MOI.OPTIMAL

        sol = result.sol
        toposol = Topology(sol.nodes, topo.arcs)
        L = pipelengths(toposol)
        c = [d.cost for d in diams]
        obj = L' * sol.lsol * c
        @assert length(obj) == 1
        @test result.value ≈ obj[1] atol=0.001
    end
end


text:
using JetPack, Jets, Test

@testset "JopLaplacian, dot product T=$(T)" for T in (Float32, Float64)
    sp = JetSpace(T,64,128)
    A = JopLaplacian(sp)
    lhs, rhs = dot_product_test(A, rand(domain(A)), rand(range(A)))
    @test lhs ≈ rhs

    sp = JetSpace(T,16,32,64)
    A = JopLaplacian(sp)
    lhs, rhs = dot_product_test(A, rand(domain(A)), rand(range(A)))
    @test lhs ≈ rhs
end

# note this is a second order derivative, so we loosen up tolerance from
# 2nd to 4th root epsilon(type)
@testset "JopLaplacian, action T=$(T)" for T in (Float32, Float64)
    tol = eps(T)^0.25

    sp = JetSpace(T,64,128)
    A = JopLaplacian(sp)
    a,b = rand(T),rand(T)
    x = zeros(sp)
    nz,nx=size(x)
    for ix=1:nx, iz=1:nz
        x[iz,ix] = a * iz^2 + b * ix^2
    end
    y = A * x
    expected = 2 * (a + b)
    z = y[2:nz-1,2:nx-1] .- expected .* ones(T,nz-2,nx-2)
    @test maximum(abs,z) < tol

    sp = JetSpace(T,16,32,64)
    A = JopLaplacian(sp)
    a,b,c = rand(T),rand(T),rand(T)
    x = zeros(sp)
    nz,ny,nx=size(x)
    for ix=1:nx, iy=1:ny, iz=1:nz
        x[iz,iy,ix] = a * iz^2 + b * iy^2 + c * iz^2
    end
    y = A * x
    expected = 2 * (a + b + c)
    z = y[2:nz-1,2:ny-1,2:nx-1] .- expected .* ones(T,nz-2,ny-2,nx-2)
    @test maximum(abs,z) < tol
end


text:
###
### UnitVariableRef
###

struct _UnitVariable{U<:Unitful.Units} <: JuMP.AbstractVariable
    variable::JuMP.ScalarVariable
    unit::U
end

"""
    UnitVariableRef(::JuMP.VariableRef, ::Unitful.Units)

A type that wraps a `VariableRef` with a `Uniful.Units`.
"""
struct UnitVariableRef{U<:Unitful.Units} <: JuMP.AbstractVariableRef
    variable::JuMP.VariableRef
    unit::U
end

JuMP.owner_model(x::UnitVariableRef) = JuMP.owner_model(x.variable)

Unitful.unit(x::UnitVariableRef) = x.unit

Base.show(io::IO, x::UnitVariableRef) = print(io, "$(x.variable) [$(x.unit)]")

function JuMP.build_variable(
    ::Function,
    info::JuMP.VariableInfo,
    unit::U,
) where {U<:Unitful.Units}
    return _UnitVariable{U}(JuMP.ScalarVariable(info), unit)
end

function JuMP.add_variable(
    model::Model,
    x::_UnitVariable{U},
    name::String,
) where {U}
    variable = JuMP.add_variable(model, x.variable, name)
    return UnitVariableRef{U}(variable, x.unit)
end

function JuMP.value(x::UnitVariableRef)
    return Unitful.Quantity(JuMP.value(x.variable), x.unit)
end

JuMP.name(x::UnitVariableRef) = JuMP.name(x.variable)

function JuMP.set_name(x::UnitVariableRef, s::String)
    return JuMP.set_name(x.variable, s)
end

function JuMP.set_start_value(x::UnitVariableRef, value::Unitful.Quantity)
    val = Unitful.uconvert(x.unit, value)
    return JuMP.set_start_value(x.variable, Unitful.ustrip(val))
end

function JuMP.set_start_value(x::UnitVariableRef, value::Real)
    return error("Start value for unit variable must be given with units")
end

function JuMP.start_value(x::UnitVariableRef)
    return Unitful.Quantity(JuMP.start_value(x.variable), x.unit)
end

JuMP.has_lower_bound(x::UnitVariableRef) = JuMP.has_lower_bound(x.variable)

function JuMP.lower_bound(x::UnitVariableRef)
    return Unitful.Quantity(JuMP.lower_bound(x.variable), x.unit)
end

function JuMP.set_lower_bound(x::UnitVariableRef, lower::Unitful.Quantity)
    val = Unitful.uconvert(x.unit, lower)
    return JuMP.set_lower_bound(x.variable, Unitful.ustrip(val))
end

function JuMP.set_lower_bound(x::UnitVariableRef, lower::Real)
    return error("Lower bound of unit variable must be given with units")
end

function JuMP.delete_lower_bound(x::UnitVariableRef)
    return JuMP.delete_lower_bound(x.variable)
end

JuMP.LowerBoundRef(x::UnitVariableRef) = JuMP.LowerBoundRef(x.variable)

JuMP.has_upper_bound(x::UnitVariableRef) = JuMP.has_upper_bound(x.variable)

function JuMP.upper_bound(x::UnitVariableRef)
    return Unitful.Quantity(JuMP.upper_bound(x.variable), x.unit)
end

function JuMP.set_upper_bound(x::UnitVariableRef, upper::Unitful.Quantity)
    val = Unitful.uconvert(x.unit, upper)
    return JuMP.set_upper_bound(x.variable, Unitful.ustrip(val))
end

function JuMP.set_upper_bound(x::UnitVariableRef, upper::Real)
    return error("Upper bound of unit variable must be given with units")
end

function JuMP.delete_upper_bound(x::UnitVariableRef)
    return JuMP.delete_upper_bound(x.variable)
end

JuMP.UpperBoundRef(x::UnitVariableRef) = JuMP.UpperBoundRef(x.variable)

JuMP.is_fixed(x::UnitVariableRef) = JuMP.is_fixed(x.variable)

function JuMP.fix_value(x::UnitVariableRef)
    return Unitful.Quantity(JuMP.fix_value(x.variable), x.unit)
end

function JuMP.fix(x::UnitVariableRef, value::Unitful.Quantity)
    val = Unitful.uconvert(x.unit, value)
    return fix(x.variable, Unitful.ustrip(val))
end

function JuMP.fix(x::UnitVariableRef, value::Real)
    return error("Fix value of unit value must be given with units")
end

JuMP.unfix(x::UnitVariableRef) = JuMP.unfix(x.variable)

JuMP.FixRef(x::UnitVariableRef) = JuMP.FixRef(x.variable)

JuMP.is_integer(x::UnitVariableRef) = JuMP.is_integer(x.variable)

JuMP.set_integer(x::UnitVariableRef) = JuMP.set_integer(x.variable)

JuMP.unset_integer(x::UnitVariableRef) = JuMP.unset_integer(x.variable)

JuMP.IntegerRef(x::UnitVariableRef) = JuMP.IntegerRef(x.variable)

JuMP.is_binary(x::UnitVariableRef) = JuMP.is_binary(x.variable)

JuMP.set_binary(x::UnitVariableRef) = JuMP.set_binary(x.variable)

JuMP.unset_binary(x::UnitVariableRef) = JuMP.unset_binary(x.variable)

JuMP.BinaryRef(x::UnitVariableRef) = JuMP.BinaryRef(x.variable)

###
### UnitExpression
###

"""
    UnitExpression(::JuMP.AbstractJuMPScalar, ::Unitful.Units)

A type that wraps an `AbstractJuMPScalar` with a `Uniful.Units`.
"""
struct UnitExpression{A<:JuMP.AbstractJuMPScalar,U<:Unitful.Units} <:
       JuMP.AbstractJuMPScalar
    expr::A
    unit::U
end

Base.show(io::IO, x::UnitExpression) = print(io, "$(x.expr) [$(x.unit)]")

Unitful.unit(x::UnitExpression) = x.unit

function Base.:(==)(x::UnitExpression, y::UnitExpression)
    return x.expr == y.expr && x.unit == y.unit
end

JuMP.moi_function(x::UnitExpression) = JuMP.moi_function(x.expr)

function JuMP.check_belongs_to_model(x::UnitExpression, model::AbstractModel)
    return JuMP.check_belongs_to_model(x.expr, model)
end

Unitful.uconvert(::U, x::UnitExpression{A,U}) where {A,U} = x

function Unitful.uconvert(unit::Unitful.Units, x::UnitExpression)
    factor = Unitful.ustrip(Unitful.uconvert(unit, Unitful.Quantity(1, x.unit)))
    return UnitExpression(factor * x.expr, unit)
end

JuMP.value(x::UnitExpression) = Unitful.Quantity(JuMP.value(x.expr), x.unit)

###
### UnitConstraintRef
###

struct _UnitConstraint{U<:Unitful.Units} <: AbstractConstraint
    constraint::ScalarConstraint
    unit::U
end

"""
    UnitConstraintRef(::JuMP.ConstraintRef, ::Unitful.Units)

A type that wraps an `ConstraintRef` with a `Uniful.Units`.
"""
struct UnitConstraintRef{U<:Unitful.Units}
    constraint::ConstraintRef
    unit::U
end

Unitful.unit(c::UnitConstraintRef) = c.unit

function Base.show(io::IO, c::UnitConstraintRef)
    return print(io, "$(c.constraint) [$(c.unit)]")
end

function JuMP.check_belongs_to_model(c::_UnitConstraint, model::AbstractModel)
    return JuMP.check_belongs_to_model(c.constraint, model)
end

function JuMP.dual(c::UnitConstraintRef)
    return Unitful.Quantity(JuMP.dual(c.constraint), Unitful.unit(1 / c.unit))
end

function JuMP.shadow_price(c::UnitConstraintRef)
    return Unitful.Quantity(
        JuMP.shadow_price(c.constraint),
        Unitful.unit(1 / c.unit),
    )
end

function JuMP.build_constraint(
    _error::Function,
    expr::UnitExpression{A,U},
    set::MOI.AbstractScalarSet,
) where {A,U}
    return _UnitConstraint{U}(
        JuMP.build_constraint(_error, expr.expr, set),
        expr.unit,
    )
end

function JuMP.build_constraint(
    _error::Function,
    uexpr::UnitExpression,
    set::MOI.AbstractScalarSet,
    unit::Unitful.Units,
)
    return JuMP.build_constraint(_error, Unitful.uconvert(unit, uexpr), set)
end

function JuMP.add_constraint(
    model::Model,
    c::_UnitConstraint{U},
    name::String,
) where {U}
    constraint = JuMP.add_constraint(model, c.constraint, name)
    return UnitConstraintRef{U}(constraint, c.unit)
end