text:
using Revise # reduce recompilation time
using LinearAlgebra
using SparseArrays
using BenchmarkTools
using UnPack
using ToeplitzMatrices
using Test
using DelimitedFiles

using KernelAbstractions
using CUDA
using CUDAapi
CUDA.allowscalar(false)

push!(LOAD_PATH, "./src")
using CommonUtils
using Basis1D
using Basis2DTri
using UniformTriMesh

push!(LOAD_PATH, "./examples/EntropyStableEuler")
using EntropyStableEuler

using SetupDG

S_N(x) = @. sin(pi*x/h)/(2*pi/h)/tan(x/2)
"""
Vandermonde matrix of sinc basis functions determined by h,
evaluated at r
"""
function vandermonde_Sinc(h,r)
    N = convert(Int, 2*pi/h)
    V = zeros(length(r),N)
    for n = 1:N
        V[:,n] = S_N(r.-n*h)
    end
    V[1,1] = 1
    V[end,end] = 1
    return V
end

vector_norm(U) = CuArrays.sum((x->x.^2).(U))

"Constants"
const sp_tol      = 1e-12
const shmem_limit = 1024  # Max number of elements in shared mem
const enable_test = false
const errTol      = 1e-5  # Error tolerance in time stepping
const gamma       = 1.4f0
const NUM_TYPE    = Float64
"Program parameters"
const compute_L2_err          = false
const add_LF_dissipation      = true

"Approximation Parameters"
const N_P   = 2;    # The order of approximation in polynomial dimension
const Np_P  = Int((N_P+1)*(N_P+2)/2)
const Np_F  = 8;    # The order of approximation in Fourier dimension
const K1D   = 60;   # Number of elements in polynomial (x,y) dimension
const Nd    = 5;
const CFL   = 0.5;
const T     = 2.5;  # End time

"Time integration Parameters"
rk4a,rk4b,rk4c = rk45_coeffs()
CN             = (N_P+1)*(N_P+2)*3/2  # estimated trace constant for CFL
dt             = CFL * 2 / CN / K1D
Nsteps         = convert(Int,ceil(T/dt))
dt             = T/Nsteps

"Initialize Reference Element in Fourier dimension"
h       = 2*pi/Np_F
column  = [0; .5*(-1).^(1:Np_F-1).*cot.((1:Np_F-1)*h/2)]
column2 = [-pi^2/3/h^2-1/6; -((-1).^(1:Np_F-1)./(2*(sin.((1:Np_F-1)*h/2)).^2))]
Dt      = Array{Float64,2}(Toeplitz(column,column[[1;Np_F:-1:2]]))
D2t     = Array{Float64,2}(Toeplitz(column2,column2[[1;Np_F:-1:2]]))
t       = LinRange(h,2*pi,Np_F)

"Initialize Reference Element in polynomial dimension"
rd = init_reference_tri(N_P);
@unpack fv,Nfaces,r,s,VDM,V1,Dr,Ds,rf,sf,wf,nrJ,nsJ,rq,sq,wq,Vq,M,Pq,Vf,LIFT = rd
Lq    = LIFT

"Problem parameters"
Nq_P   = length(rq)
Nfp_P  = length(rf)
Nh_P   = Nq_P+Nfp_P # Number of hybridized points
Nq     = Nq_P*Np_F
Nh     = Nh_P*Np_F
Nfp    = Nfp_P*Np_F

# Adpatively determine number of threads
thread_count = 2^8
for n = 8:-1:5
    global thread_count
    thread_count = 2^n
    Nk_tri_max = div(thread_count-2,Nh_P)+2 # Max amount of triangles in a block
    Nk_fourier_max = div(thread_count-2,Np_F)+2 # Max amount of fourier slides in a block
    if max(Nd*Nh_P*Nk_tri_max,Nd*Np_F*Nk_fourier_max) < shmem_limit
        break
    end
end
const num_threads = thread_count
Nk_max = num_threads == 1 ? 1 : div(num_threads-2,Nh)+2 # Max amount of element in a block
Nk_tri_max = num_threads == 1 ? 1 : div(num_threads-2,Nh_P)+2 # Max amount of triangles in a block
Nk_fourier_max = num_threads == 1 ? 1 : div(num_threads-2,Np_F)+2 # Max amount of fourier slides in a block

"Mesh related variables"
# Initialize 2D triangular mesh
VX,VY,EToV = uniform_tri_mesh(K1D,K1D)
@. VX = VX
@. VY = VY
md   = init_mesh((VX,VY),EToV,rd)
# Intialize triangular prism
VX   = repeat(VX,2)
VY   = repeat(VY,2)
VZ   = [-1 .+ 2/Np_F*ones((K1D+1)*(K1D+1),1); ones((K1D+1)*(K1D+1),1)]
EToV = [EToV EToV.+(K1D+1)*(K1D+1)]

# Make domain periodic
@unpack Nfaces,Vf = rd
@unpack xf,yf,K,mapM,mapP,mapB = md
LX,LY = (x->maximum(x)-minimum(x)).((VX,VY)) # find lengths of domain
mapPB = build_periodic_boundary_maps(xf,yf,LX,LY,Nfaces*K,mapM,mapP,mapB)
mapP[mapB] = mapPB

# Initialize 3D mesh
@unpack x,y,xf,yf,xq,yq,rxJ,sxJ,ryJ,syJ,J,nxJ,nyJ,sJ,mapM,mapP,mapB = md
x,y,xf,yf,xq,yq,rxJ,sxJ,ryJ,syJ,J,nxJ,nyJ,sJ = (x->reshape(repeat(x,inner=(1,Np_F)),size(x,1),Np_F,K)).((x,y,xf,yf,xq,yq,rxJ,sxJ,ryJ,syJ,J,nxJ,nyJ,sJ))
z,zq,zf = (x->reshape(repeat(collect(-1+2/Np_F:(2/Np_F):1),inner=(1,x),outer=(K,1))',x,Np_F,K)).((Np_P,Nq_P,Nfp_P))
mapM = reshape(1:Nfp_P*Np_P*K,Nfp_P,Np_P,K)
mapP_2D = (x->mod1(x,Nfp_P)+div(x-1,Nfp_P)*Nfp_P*Np_F).(mapP)
mapP = reshape(repeat(mapP_2D,inner=(1,Np_F)),Nfp_P,Np_F,K)
for j = 1:Np_F
    mapP[:,j,:] = mapP[:,j,:].+(j-1)*Nfp_P
end
mapP_d = zeros(Int32,Nfp_P,Np_F,Nd,K)
for k = 1:K
    for j = 1:Np_F
        for i = 1:Nfp_P
            val = mapP[i,j,k]
            elem = div(val-1,Nfp)
            n = mod1(val,Nfp)
            mapP_d[i,j,1,k] = elem*Nd*Nfp+n
            mapP_d[i,j,2,k] = elem*Nd*Nfp+n+Nfp
            mapP_d[i,j,3,k] = elem*Nd*Nfp+n+2*Nfp
            mapP_d[i,j,4,k] = elem*Nd*Nfp+n+3*Nfp
            mapP_d[i,j,5,k] = elem*Nd*Nfp+n+4*Nfp
        end
    end
end
mapP = mapP[:]
mapP_d = mapP_d[:]

#############################################
######  Construct hybridized operators ######
#############################################

# scale by Fourier dimension
M = h*M
wq = h*wq
wf = h*wf

# Hybridized operators
Vh = [Vq;Vf]
rxJ,sxJ,ryJ,syJ = (x->mapslices((y->Vh*y),x,dims=(1,2))).((rxJ,sxJ,ryJ,syJ))
Ef = Vf*Pq
Br,Bs = (x->diagm(wf.*x)).((nrJ,nsJ))
Qr,Qs = (x->Pq'*M*x*Pq).((Dr,Ds))
Qrh,Qsh = (x->1/2*[x[1]-x[1]' Ef'*x[2];
                   -x[2]*Ef   x[2]]).(((Qr,Br),(Qs,Bs)))
Qrh_skew,Qsh_skew = (x->1/2*(x-x')).((Qrh,Qsh))
Qt = Dt
Qth = Qt # Not the SBP operator, weighted when flux differencing
Ph = [Vq;Vf]*Pq

# TODO: assume mesh uniform affine, so Jacobian are constants
# TODO: fix other Jacobian parts
JP = 1/K1D^2
JF = 1/pi
J = JF*JP
wq = J*wq
wf = JF*wf
Lq = 1/JP*Lq
Qrh = JF*Qrh
Qsh = JF*Qsh
Qth = JP*Qth
Qrh_skew = 1/2*(Qrh-Qrh')
Qsh_skew = 1/2*(Qsh-Qsh')
LIFTq = Vq*Lq
VPh = Vq*[Pq Lq]*diagm(1 ./ [wq;wf])
Winv = diagm(1 ./ [wq;wf])
Wq = diagm(wq)

# Assume affine meshes
rxJ = [rxJ[1,1,k] for k in 1:K]
sxJ = [sxJ[1,1,k] for k in 1:K]
ryJ = [ryJ[1,1,k] for k in 1:K]
syJ = [syJ[1,1,k] for k in 1:K]

# convert precision
Vq,wq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq,VPh,Wq,rxJ,sxJ,ryJ,syJ,nxJ,nyJ,sJ,J,h,rk4a,rk4b = (A->convert.(NUM_TYPE,A)).((Vq,wq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq,VPh,Wq,rxJ,sxJ,ryJ,syJ,nxJ,nyJ,sJ,J,h,rk4a,rk4b))

# TODO: refactor
Vq,Vf,wq,wf,Pq,Lq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq,VPh,Wq,Winv,rxJ,sxJ,ryJ,syJ = (x->CuArray(x)).((Vq,Vf,wq,wf,Pq,Lq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq,VPh,Wq,Winv,rxJ,sxJ,ryJ,syJ))
nxJ = CuArray(nxJ[:])
nyJ = CuArray(nyJ[:])
sJ  = CuArray(sJ[:])

ops = (Vq,wq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq,VPh,Wq)
mesh = (rxJ,sxJ,ryJ,syJ,nxJ,nyJ,sJ,J,h,mapP,mapP_d)
param = (K,Np_P,Nq_P,Nfp_P,Nh_P,Np_F,Nq,Nfp,Nh,Nk_max,Nk_tri_max,Nk_fourier_max)

# ================================= #
# ============ Routines =========== #
# ================================= #
function u_to_v!(VU,Q,Nd,Nq)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x

    @inbounds begin
    for i = index:stride:ceil(Int,length(VU)/Nd)
        k = div(i-1,Nq) # Current element
        n = mod1(i,Nq) # Current quad node

        idx = k*Nd*Nq+n
        rho = Q[idx]
        rhou = Q[idx+Nq]
        rhov = Q[idx+2*Nq]
        rhow = Q[idx+3*Nq]
        E = Q[idx+4*Nq]
        rhoUnorm = rhou^2+rhov^2+rhow^2
        rhoe = E-.5*rhoUnorm/rho
        sU = CUDA.log((gamma-1.0f0)*rhoe/CUDA.exp(gamma*CUDA.log(rho)))

        VU[idx] = (-E+rhoe*(gamma+1.0f0-sU))/rhoe
        VU[idx+Nq] = rhou/rhoe
        VU[idx+2*Nq] = rhov/rhoe
        VU[idx+3*Nq] = rhow/rhoe
        VU[idx+4*Nq] = -rho/rhoe
    end
    end
end

function v_to_u!(Qh,Nd,Nh)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x

    @inbounds begin
    for i = index:stride:ceil(Int,length(Qh)/Nd)
        k = div(i-1,Nh) # Current element
        n = mod1(i,Nh) # Current hybridized node

        idx = k*Nd*Nh+n
        v1 = Qh[idx     ]
        vu = Qh[idx+  Nh]
        vv = Qh[idx+2*Nh]
        vw = Qh[idx+3*Nh]
        vE = Qh[idx+4*Nh]
        vUnorm = vu^2+vv^2+vw^2
        rhoeV = CUDA.exp(1.0f0/(gamma-1.0f0)*CUDA.log((gamma-1.0f0)/CUDA.exp(gamma*CUDA.log(-vE))))*CUDA.exp(-(gamma-v1+vUnorm/(2.0f0*vE))/(gamma-1.0f0))

        Qh[idx     ] = -rhoeV*vE
        Qh[idx+  Nh] = rhoeV*vu
        Qh[idx+2*Nh] = rhoeV*vv
        Qh[idx+3*Nh] = rhoeV*vw
        Qh[idx+4*Nh] = rhoeV*(1.0f0-vUnorm/(2.0f0*vE))
    end
    end
end

#TODO: combine with v_to_u
function u_to_primitive!(Qh,Nd,Nh)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x

    @inbounds begin
    for i = index:stride:ceil(Int,length(Qh)/Nd)
        k = div(i-1,Nh) # Current element
        n = mod1(i,Nh) # Current hybridized node
        idx = k*Nd*Nh+n

        rho = Qh[idx]
        rhou = Qh[idx+Nh]
        rhov = Qh[idx+2*Nh]
        rhow = Qh[idx+3*Nh]
        E = Qh[idx+4*Nh]

        beta = rho/(2.0f0*(gamma-1.0f0)*(E-.5*(rhou^2+rhov^2+rhow^2)/rho))

        Qh[idx+Nh] = rhou/rho
        Qh[idx+2*Nh] = rhov/rho
        Qh[idx+3*Nh] = rhow/rho
        Qh[idx+4*Nh] = beta
    end
    end
end

function extract_face_val_conservative!(Uf,Qh,Nfp_P,Nq_P,Nh_P)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x

    @inbounds begin
    for i = index:stride:length(Uf)
        k = div(i-1,Nfp_P)
        n = mod1(i,Nfp_P)
        Uf[i] = Qh[k*Nh_P+Nq_P+n]
    end
    end
end

function extract_face_val!(QM,Qh,Nfp_P,Nq_P,Nh_P)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x

    @inbounds begin
    for i = index:stride:length(QM)
        k = div(i-1,Nfp_P)
        n = mod1(i,Nfp_P)
        QM[i] = Qh[k*Nh_P+Nq_P+n]
    end
    end
end

function CU_logmean(uL,uR,logL,logR)
    da = uR-uL
    aavg = .5f0*(uL+uR)
    f = da/aavg
    if CUDA.abs(f)<1e-4
        v = f^2
        return aavg*(1.0f0 + v*(-.2f0-v*(.0512f0 - v*0.026038857142857f0)))
    else
        return -da/(logL-logR)
    end
end

function CU_euler_flux(rhoM,uM,vM,wM,betaM,rhoP,uP,vP,wP,betaP,rhologM,betalogM,rhologP,betalogP)
    rholog  = CU_logmean(rhoM,rhoP,rhologM,rhologP)
    betalog = CU_logmean(betaM,betaP,betalogM,betalogP)

    rhoavg = .5f0*(rhoM+rhoP)
    uavg   = .5f0*(uM+uP)
    vavg   = .5f0*(vM+vP)
    wavg   = .5f0*(wM+wP)

    unorm = uM*uP+vM*vP+wM*wP
    pa = rhoavg/(betaM+betaP)
    E_plus_p = rholog/(2.0f0*(gamma-1.0f0)*betalog) + pa + .5f0*rholog*unorm

    FxS1 = (@. rholog*uavg)
    FxS2 = (@. FxS1*uavg + pa)
    FxS3 = (@. FxS1*vavg) # rho * u * v
    FxS4 = (@. FxS1*wavg) # rho * u * w
    FxS5 = (@. E_plus_p*uavg)
    FyS1 = (@. rholog*vavg)
    FyS2 = (@. FxS3) # rho * u * v
    FyS3 = (@. FyS1*vavg + pa)
    FyS4 = (@. FyS1*wavg) # rho * v * w
    FyS5 = (@. E_plus_p*vavg)
    FzS1 = (@. rholog*wavg)
    FzS2 = (@. FxS4) # rho * u * w
    FzS3 = (@. FyS4) # rho * v * w
    FzS4 = (@. FzS1*wavg + pa) # rho * w^2 + p
    FzS5 = (@. E_plus_p*wavg)

    return FxS1,FxS2,FxS3,FxS4,FxS5,FyS1,FyS2,FyS3,FyS4,FyS5,FzS1,FzS2,FzS3,FzS4,FzS5
end

function construct_lam!(lam,Uf,nxJ,nyJ,sJ,Nfp,K,Nd)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x

    @inbounds begin
    for i = index:stride:Nfp*K
        k = div(i-1,Nfp)
        n = mod1(i,Nfp)

        rhoM  = Uf[k*Nd*Nfp+n      ]
        rhouM = Uf[k*Nd*Nfp+n+  Nfp]
        rhovM = Uf[k*Nd*Nfp+n+2*Nfp]
        rhowM = Uf[k*Nd*Nfp+n+3*Nfp]
        EM    = Uf[k*Nd*Nfp+n+4*Nfp]
        nxJ_val = nxJ[k*Nfp+n]
        nyJ_val = nyJ[k*Nfp+n]
        sJ_val  = sJ[k*Nfp+n]
        rhouM_n = (rhouM*nxJ_val+rhovM*nyJ_val)/sJ_val
        lam[k*Nfp+n] = CUDA.sqrt(CUDA.abs(rhouM_n/rhoM))+CUDA.sqrt(gamma*(gamma-1.0f0)*(EM-0.5f0*rhouM_n^2/rhoM)/rhoM)
    end
    end
end

function surface_kernel!(flux,QM,QP,Uf,UfP,LFc,nxJ,nyJ,Nfp,K,Nd)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x
    @inbounds begin
    for i = index:stride:Nfp*K
        k = div(i-1,Nfp)
        n = mod1(i,Nfp)

        rhoM  = QM[k*Nd*Nfp+n      ]
        uM    = QM[k*Nd*Nfp+n+  Nfp]
        vM    = QM[k*Nd*Nfp+n+2*Nfp]
        wM    = QM[k*Nd*Nfp+n+3*Nfp]
        betaM = QM[k*Nd*Nfp+n+4*Nfp]
        rhoP  = QP[k*Nd*Nfp+n      ]
        uP    = QP[k*Nd*Nfp+n+  Nfp]
        vP    = QP[k*Nd*Nfp+n+2*Nfp]
        wP    = QP[k*Nd*Nfp+n+3*Nfp]
        betaP = QP[k*Nd*Nfp+n+4*Nfp]
        nxJ_val = nxJ[k*Nfp+n]
        nyJ_val = nyJ[k*Nfp+n]

        rhouM = Uf[k*Nd*Nfp+n+  Nfp]
        rhovM = Uf[k*Nd*Nfp+n+2*Nfp]
        rhowM = Uf[k*Nd*Nfp+n+3*Nfp]
        EM    = Uf[k*Nd*Nfp+n+4*Nfp]
        rhouP = UfP[k*Nd*Nfp+n+  Nfp]
        rhovP = UfP[k*Nd*Nfp+n+2*Nfp]
        rhowP = UfP[k*Nd*Nfp+n+3*Nfp]
        EP    = UfP[k*Nd*Nfp+n+4*Nfp]
        LFc_val = LFc[k*Nfp+n]

        rhologM = CUDA.log(rhoM)
        rhologP = CUDA.log(rhoP)
        betalogM = CUDA.log(betaM)
        betalogP = CUDA.log(betaP)

        FxS1,FxS2,FxS3,FxS4,FxS5,FyS1,FyS2,FyS3,FyS4,FyS5,_,_,_,_,_ = CU_euler_flux(rhoM,uM,vM,wM,betaM,rhoP,uP,vP,wP,betaP,rhologM,betalogM,rhologP,betalogP)

        flux[k*Nd*Nfp+n      ] = nxJ_val*FxS1+nyJ_val*FyS1-LFc_val*(rhoP-rhoM)
        flux[k*Nd*Nfp+n+  Nfp] = nxJ_val*FxS2+nyJ_val*FyS2-LFc_val*(rhouP-rhouM)
        flux[k*Nd*Nfp+n+2*Nfp] = nxJ_val*FxS3+nyJ_val*FyS3-LFc_val*(rhovP-rhovM)
        flux[k*Nd*Nfp+n+3*Nfp] = nxJ_val*FxS4+nyJ_val*FyS4-LFc_val*(rhowP-rhowM)
        flux[k*Nd*Nfp+n+4*Nfp] = nxJ_val*FxS5+nyJ_val*FyS5-LFc_val*(EP-EM)
    end
    end
end

function flux_differencing_xy_kernel!(gradfh,Qh,rxJ,sxJ,ryJ,syJ,Qrh_skew,Qsh_skew,Nh,Nh_P,Np_F,K,Nd,Nq_P,num_threads,Nk_max,Nk_tri_max)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x
    #TODO: using stride to avoid possible insufficient blocks?
    tid_start = (blockIdx().x-1)*blockDim().x + 1 # starting threadid of current block
    tid_end = (blockIdx().x-1)*blockDim().x + num_threads # ending threadid of current block
    k_start = div(tid_start-1,Nh) # starting elementid of current block
    k_end = tid_end < Nh*K ? div(tid_end-1,Nh) : K-1 # ending elementid of current block
    Nk = k_end-k_start+1 # Number of elements in current block
    k_tri_start = div(tid_start-1,Nh_P) # starting triangle id of current block
    k_tri_end = tid_end < Nh*K ? div(tid_end-1,Nh_P) : Np_F*K-1 # ending triangle id of current block
    Nk_tri = k_tri_end-k_tri_start+1 # Number of triangles in current block

    # Parallel read
    Qh_shared = @cuDynamicSharedMem(NUM_TYPE,Nh_P*Nd*Nk_tri_max)
    load_size = div(Nh_P*Nd*Nk_tri-1,num_threads)+1 # size of read of each thread
    @inbounds begin
    for i in (threadIdx().x-1)*load_size+1:min(threadIdx().x*load_size,Nh_P*Nd*Nk_tri)
        offset_tri = div(i-1,Nh_P*Nd) # distance with k_tri_start
        n_local = mod1(i,Nh_P*Nd)
        n_d = div(n_local-1,Nh_P) # Current component
        m_xy = mod1(n_local,Nh_P) # Current hybridized node index at xy slice
        i_tri = k_tri_start+offset_tri # Current triangle
        k = div(i_tri,Np_F) # Current element
        i_xy = mod(i_tri,Np_F) # Current local xy slice id on the elment
        Qh_shared[i] = Qh[k*Nd*Nh+i_xy*Nh_P+m_xy+n_d*Nh]
    end
    end

    sync_threads()

    @inbounds begin
    if index <= Nh*K
        i = index
        k = div(i-1,Nh)   # Current element
        m = mod1(i,Nh)    # Current node on element
        i_xy = div(m-1,Nh_P) # Current local xy slice id on the element
        m_xy = mod1(i,Nh_P)  # Current hybridized node index at x-y slice

        k_tri_offset = div(i-1,Nh_P)-k_tri_start
        rhoL  = Qh_shared[k_tri_offset*Nd*Nh_P+m_xy       ]
        uL    = Qh_shared[k_tri_offset*Nd*Nh_P+m_xy+  Nh_P]
        vL    = Qh_shared[k_tri_offset*Nd*Nh_P+m_xy+2*Nh_P]
        wL    = Qh_shared[k_tri_offset*Nd*Nh_P+m_xy+3*Nh_P]
        betaL = Qh_shared[k_tri_offset*Nd*Nh_P+m_xy+4*Nh_P]
        rhologL  = CUDA.log(rhoL)
        betalogL = CUDA.log(betaL)

        rho_sum = 0.0f0
        u_sum = 0.0f0
        v_sum = 0.0f0
        w_sum = 0.0f0
        beta_sum = 0.0f0

        # Assume Affine meshes

        rxJ_val = rxJ[k+1]
        sxJ_val = sxJ[k+1]
        ryJ_val = ryJ[k+1]
        syJ_val = syJ[k+1]

        for n_xy = 1:Nh_P
            if m_xy <= Nq_P || n_xy <= Nq_P
                rhoR  = Qh_shared[k_tri_offset*Nd*Nh_P+n_xy       ]
                uR    = Qh_shared[k_tri_offset*Nd*Nh_P+n_xy+  Nh_P]
                vR    = Qh_shared[k_tri_offset*Nd*Nh_P+n_xy+2*Nh_P]
                wR    = Qh_shared[k_tri_offset*Nd*Nh_P+n_xy+3*Nh_P]
                betaR = Qh_shared[k_tri_offset*Nd*Nh_P+n_xy+4*Nh_P]
                rhologR = CUDA.log(rhoR)
                betalogR = CUDA.log(betaR)

                FxS1,FxS2,FxS3,FxS4,FxS5,FyS1,FyS2,FyS3,FyS4,FyS5,FzS1,FzS2,FzS3,FzS4,FzS5 = CU_euler_flux(rhoL,uL,vL,wL,betaL,rhoR,uR,vR,wR,betaR,rhologL,betalogL,rhologR,betalogR)

                Qx_val = 2.0f0*(rxJ_val*Qrh_skew[m_xy,n_xy]+sxJ_val*Qsh_skew[m_xy,n_xy])
                Qy_val = 2.0f0*(ryJ_val*Qrh_skew[m_xy,n_xy]+syJ_val*Qsh_skew[m_xy,n_xy])
                rho_sum  += Qx_val*FxS1+Qy_val*FyS1
                u_sum    += Qx_val*FxS2+Qy_val*FyS2
                v_sum    += Qx_val*FxS3+Qy_val*FyS3
                w_sum    += Qx_val*FxS4+Qy_val*FyS4
                beta_sum += Qx_val*FxS5+Qy_val*FyS5
            end
        end

        gradfh[k*Nd*Nh+m     ] = rho_sum
        gradfh[k*Nd*Nh+m+  Nh] = u_sum
        gradfh[k*Nd*Nh+m+2*Nh] = v_sum
        gradfh[k*Nd*Nh+m+3*Nh] = w_sum
        gradfh[k*Nd*Nh+m+4*Nh] = beta_sum
    end
    end
    return
end

function flux_differencing_z_kernel!(gradfh,Qh,wq,h,J,Qth,Nh,Nq,Nh_P,Np_F,K,Nd,Nq_P,num_threads,Nk_max,Nk_fourier_max)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x
    #TODO: using stride to avoid possible insufficient blocks?
    tid_start = (blockIdx().x-1)*blockDim().x + 1 # starting threadid of current block
    tid_end = (blockIdx().x-1)*blockDim().x + num_threads # ending threadid of current block
    k_start = div(tid_start-1,Nq) # starting elementid of current block
    k_end = tid_end < Nq*K ? div(tid_end-1,Nq) : K-1 # ending elementid of current block
    Nk = k_end-k_start+1 # Number of elements in current block
    k_fourier_start = div(tid_start-1,Np_F)
    k_fourier_end = tid_end < Nq*K ? div(tid_end-1,Np_F) : Nq_P*K-1
    Nk_fourier = k_fourier_end-k_fourier_start+1

    Qh_shared = @cuDynamicSharedMem(NUM_TYPE,Np_F*Nd*Nk_fourier_max)
    load_size = div(Np_F*Nd*Nk_fourier-1,num_threads)+1 # size of read of each thread
    @inbounds begin
    for i in (threadIdx().x-1)*load_size+1:min(threadIdx().x*load_size,Np_F*Nd*Nk_fourier)
        offset_fourier = div(i-1,Np_F*Nd) # distance with k_fourier_start
        n_local = mod1(i,Np_F*Nd)
        n_d = div(n_local-1,Np_F) # Current component
        m_z = mod1(n_local,Np_F) # Current local node index at fourier slice
        i_fourier = k_fourier_start+offset_fourier
        k = div(i_fourier,Nq_P) # Current element
        i_z = mod(i_fourier,Nq_P) # Current local z slice id on the element
        Qh_shared[i] = Qh[k*Nd*Nh+(m_z-1)*Nh_P+i_z+1+n_d*Nh]
    end
    end
    sync_threads()

    @inbounds begin
    if index <= Nq*K
        i = index
        k = div(i-1,Nq) # Current element
        m = mod1(i,Nq) # local node number in current element
        i_z = div(m-1,Np_F) # Fourier element id on current element
        m_z = mod1(m,Np_F) # local node number on Fourier element

        k_fourier_offset = div(i-1,Np_F)-k_fourier_start
        rhoL  = Qh_shared[k_fourier_offset*Nd*Np_F+m_z       ]
        uL    = Qh_shared[k_fourier_offset*Nd*Np_F+m_z+  Np_F]
        vL    = Qh_shared[k_fourier_offset*Nd*Np_F+m_z+2*Np_F]
        wL    = Qh_shared[k_fourier_offset*Nd*Np_F+m_z+3*Np_F]
        betaL = Qh_shared[k_fourier_offset*Nd*Np_F+m_z+4*Np_F]
        rhologL = CUDA.log(rhoL)
        betalogL = CUDA.log(betaL)

        rho_sum = 0.0f0
        u_sum = 0.0f0
        v_sum = 0.0f0
        w_sum = 0.0f0
        beta_sum = 0.0f0

        for n_z = 1:Np_F
            rhoR  = Qh_shared[k_fourier_offset*Nd*Np_F+n_z       ]
            uR    = Qh_shared[k_fourier_offset*Nd*Np_F+n_z+  Np_F]
            vR    = Qh_shared[k_fourier_offset*Nd*Np_F+n_z+2*Np_F]
            wR    = Qh_shared[k_fourier_offset*Nd*Np_F+n_z+3*Np_F]
            betaR = Qh_shared[k_fourier_offset*Nd*Np_F+n_z+4*Np_F]

            rhologR = CUDA.log(rhoR)
            betalogR = CUDA.log(betaR)

            _,_,_,_,_,_,_,_,_,_,FzS1,FzS2,FzS3,FzS4,FzS5 = CU_euler_flux(rhoL,uL,vL,wL,betaL,rhoR,uR,vR,wR,betaR,rhologL,betalogL,rhologR,betalogR)

            Qz_val = 2.0f0/J*wq[i_z+1]*Qth[m_z,n_z]
            rho_sum += Qz_val*FzS1
            u_sum += Qz_val*FzS2
            v_sum += Qz_val*FzS3
            w_sum += Qz_val*FzS4
            beta_sum += Qz_val*FzS5
        end
        gradfh[k*Nd*Nh+(m_z-1)*Nh_P+i_z+1     ] += rho_sum
        gradfh[k*Nd*Nh+(m_z-1)*Nh_P+i_z+1+  Nh] += u_sum
        gradfh[k*Nd*Nh+(m_z-1)*Nh_P+i_z+1+2*Nh] += v_sum
        gradfh[k*Nd*Nh+(m_z-1)*Nh_P+i_z+1+3*Nh] += w_sum
        gradfh[k*Nd*Nh+(m_z-1)*Nh_P+i_z+1+4*Nh] += beta_sum
    end
    end
    return
end

function volume_kernel!(gradfh,Qh,rxJ,sxJ,ryJ,syJ,wq,h,J,Qrh_skew,Qsh_skew,Qth,Nh,Nh_P,Np_F,K,Nd,Nq_P,num_threads,Nk_max)
    index = (blockIdx().x-1)*blockDim().x + threadIdx().x
    stride = blockDim().x*gridDim().x
    #TODO: using stride to avoid possible insufficient blocks?
    tid_start = (blockIdx().x-1)*blockDim().x + 1 # starting threadid of current block
    tid_end = (blockIdx().x-1)*blockDim().x + num_threads # ending threadid of current block
    k_start = div(tid_start-1,Nh) # starting elementid of current block
    k_end = tid_end < Nh*K ? div(tid_end-1,Nh) : K-1 # ending elementid of current block
    Nk = k_end-k_start+1 # Number of elements in current block

    # Parallel read
    Qh_shared = @cuDynamicSharedMem(NUM_TYPE,Nh*Nd*Nk_max)
    load_size = div(Nh*Nd*Nk-1,num_threads)+1 # size of read of each thread

    @inbounds begin
    for r_id in (threadIdx().x-1)*load_size+1:min(threadIdx().x*load_size,Nh*Nd*Nk)
        Qh_shared[r_id] = Qh[k_start*Nd*Nh+r_id]
    end
    end
    sync_threads()

    @inbounds begin
    if index <= Nh*K
        i = index
        k = div(i-1,Nh)
        m = mod1(i,Nh)

        k_offset = k-k_start

        rhoL  = Qh_shared[k_offset*Nd*Nh+m     ]
        uL    = Qh_shared[k_offset*Nd*Nh+m+  Nh]
        vL    = Qh_shared[k_offset*Nd*Nh+m+2*Nh]
        wL    = Qh_shared[k_offset*Nd*Nh+m+3*Nh]
        betaL = Qh_shared[k_offset*Nd*Nh+m+4*Nh]
        rhologL = CUDA.log(rhoL)
        betalogL = CUDA.log(betaL)

        t = div(m-1,Nh_P) # Current x-y slice
        s = mod1(m,Nh_P) # Current hybridized node index at x-y slice
        xy_idx = t*Nh_P+1:(t+1)*Nh_P # Nonzero index for Qrh, Qsh
        z_idx = s:Nh_P:s+(Np_F-1)*Nh_P

        rho_sum = 0.0f0
        u_sum = 0.0f0
        v_sum = 0.0f0
        w_sum = 0.0f0
        beta_sum = 0.0f0

        # Assume Affine meshes

        rxJ_val = rxJ[k+1]
        sxJ_val = sxJ[k+1]
        ryJ_val = ryJ[k+1]
        syJ_val = syJ[k+1]

        # TODO: better way to indexing
        for n = 1:Nh
            if n in xy_idx || n in z_idx

                rhoR  = Qh_shared[k_offset*Nd*Nh+n     ]
                uR    = Qh_shared[k_offset*Nd*Nh+n+  Nh]
                vR    = Qh_shared[k_offset*Nd*Nh+n+2*Nh]
                wR    = Qh_shared[k_offset*Nd*Nh+n+3*Nh]
                betaR = Qh_shared[k_offset*Nd*Nh+n+4*Nh]
                rhologR = CUDA.log(rhoR)
                betalogR = CUDA.log(betaR)

                FxS1,FxS2,FxS3,FxS4,FxS5,FyS1,FyS2,FyS3,FyS4,FyS5,FzS1,FzS2,FzS3,FzS4,FzS5 = CU_euler_flux(rhoL,uL,vL,wL,betaL,rhoR,uR,vR,wR,betaR,rhologL,betalogL,rhologR,betalogR)

                if n in xy_idx
                    col_idx = mod1(n,Nh_P)
                    Qx_val = 2.0f0*(rxJ_val*Qrh_skew[s,col_idx]+sxJ_val*Qsh_skew[s,col_idx])
                    Qy_val = 2.0f0*(ryJ_val*Qrh_skew[s,col_idx]+syJ_val*Qsh_skew[s,col_idx])
                    rho_sum += Qx_val*FxS1+Qy_val*FyS1
                    u_sum += Qx_val*FxS2+Qy_val*FyS2
                    v_sum += Qx_val*FxS3+Qy_val*FyS3
                    w_sum += Qx_val*FxS4+Qy_val*FyS4
                    beta_sum += Qx_val*FxS5+Qy_val*FyS5
                end
                #=
                if n in z_idx && s <= Nq_P
                    col_idx = div(n-1,Nh_P)+1
                    wqn = 2.0f0/J*wq[s]
                    Qz_val = wqn*Qth[t+1,col_idx]
                    rho_sum += Qz_val*FzS1
                    u_sum += Qz_val*FzS2
                    v_sum += Qz_val*FzS3
                    w_sum += Qz_val*FzS4
                    beta_sum += Qz_val*FzS5
                end
                =#
            end
        end

        gradfh[k*Nd*Nh+m     ] = rho_sum
        gradfh[k*Nd*Nh+m+  Nh] = u_sum
        gradfh[k*Nd*Nh+m+2*Nh] = v_sum
        gradfh[k*Nd*Nh+m+3*Nh] = w_sum
        gradfh[k*Nd*Nh+m+4*Nh] = beta_sum
    end
    end
    return
end

# ================================= #
# ============ Routines =========== #
# ================================= #

function rhs(Q,ops,mesh,param,num_threads,compute_rhstest,enable_test)
    Vq,wq,Qrh_skew,Qsh_skew,Qth,Ph,LIFTq,VPh,Wq = ops
    rxJ,sxJ,ryJ,syJ,nxJ,nyJ,sJ,J,h,mapP,mapP_d = mesh
    K,Np_P,Nq_P,Nfp_P,Nh_P,Np_F,Nq,Nfp,Nh,Nk_max,Nk_tri_max,Nk_fourier_max = param

    VU = CUDA.fill(CUDA.zero(NUM_TYPE),Nq*Nd*K)
    Uf = CUDA.fill(CUDA.zero(NUM_TYPE),Nfp*Nd*K)
    QM = CUDA.fill(CUDA.zero(NUM_TYPE),Nfp*Nd*K)
    flux = CUDA.fill(CUDA.zero(NUM_TYPE),Nfp*K*Nd)
    gradfh = CUDA.fill(CUDA.zero(NUM_TYPE),Nh*Nd*K)
    lam = CUDA.fill(CUDA.zero(NUM_TYPE),Nfp*K)
    LFc = CUDA.fill(CUDA.zero(NUM_TYPE),Nfp*K)

    # Entropy Projection
    @cuda threads=num_threads blocks=ceil(Int,Nq*K/num_threads) u_to_v!(VU,Q,Nd,Nq)
    synchronize()
    Qh = reshape(Ph*reshape(VU,Nq_P,Np_F*Nd*K),Nh*Nd*K)
    @cuda threads=num_threads blocks=ceil(Int,Nh*K/num_threads) v_to_u!(Qh,Nd,Nh)
    synchronize()
    @cuda threads=num_threads blocks=ceil(Int,Nfp*Nd*K/num_threads) extract_face_val_conservative!(Uf,Qh,Nfp_P,Nq_P,Nh_P)
    UfP = Uf[mapP_d]
    synchronize()
    @cuda threads=num_threads blocks=ceil(Int,Nh*K/num_threads) u_to_primitive!(Qh,Nd,Nh)
    synchronize()

    # Compute Surface values
    @cuda threads=num_threads blocks=ceil(Int,Nfp*Nd*K/num_threads) extract_face_val!(QM,Qh,Nfp_P,Nq_P,Nh_P)
    synchronize()
    QP = QM[mapP_d]

    # LF dissipation
    @cuda threads=num_threads blocks=ceil(Int,Nfp*K/num_threads) construct_lam!(lam,Uf,nxJ,nyJ,sJ,Nfp,K,Nd)
    synchronize()
    LFc .= .5f0*CUDA.max.(lam,lam[mapP]).*sJ

    # Surface kernel
    @cuda threads=num_threads blocks=ceil(Int,Nfp*K/num_threads) surface_kernel!(flux,QM,QP,Uf,UfP,LFc,nxJ,nyJ,Nfp,K,Nd)
    synchronize()
    flux = reshape(LIFTq*reshape(flux,Nfp_P,Np_F*Nd*K),Nq*Nd*K)

    # Volume kernel
    @cuda threads=num_threads blocks=ceil(Int,Nh*K/num_threads) shmem=sizeof(NUM_TYPE)*Nh_P*Nd*Nk_tri_max flux_differencing_xy_kernel!(gradfh,Qh,rxJ,sxJ,ryJ,syJ,Qrh_skew,Qsh_skew,Nh,Nh_P,Np_F,K,Nd,Nq_P,num_threads,Nk_max,Nk_tri_max)
    @cuda threads=num_threads blocks=ceil(Int,Nq*K/num_threads) shmem=sizeof(NUM_TYPE)*Np_F*Nd*Nk_fourier_max flux_differencing_z_kernel!(gradfh,Qh,wq,h,J,Qth,Nh,Nq,Nh_P,Np_F,K,Nd,Nq_P,num_threads,Nk_max,Nk_fourier_max)
    #@cuda threads=num_threads blocks=ceil(Int,Nh*K/num_threads) shmem=sizeof(NUM_TYPE)*Nh*Nd*Nk_max volume_kernel!(gradfh,Qh,rxJ,sxJ,ryJ,syJ,wq,h,J,Qrh_skew,Qsh_skew,Qth,Nh,Nh_P,Np_F,K,Nd,Nq_P,num_threads,Nk_max)
    synchronize()
    gradf = reshape(VPh*reshape(gradfh,Nh_P,Np_F*Nd*K),Nq*Nd*K)

    # Combine
    rhsQ = -(gradf+flux)

    # Compute rhstest
    rhstest = 0
    if compute_rhstest
        rhstest = CUDA.sum(Wq*reshape(VU.*rhsQ,Nq_P,Np_F*Nd*K))
    end

    if enable_test
        @show CUDA.maximum(VU)
        @show CUDA.minimum(VU)
        @show CUDA.sum(VU)
        @show CUDA.maximum(Qh)
        @show CUDA.minimum(Qh)
        @show CUDA.sum(Qh)
        @show CUDA.maximum(QM)
        @show CUDA.minimum(QM)
        @show CUDA.sum(QM)
        @show CUDA.maximum(QP)
        @show CUDA.minimum(QP)
        @show CUDA.sum(QP)
        @show CUDA.maximum(lam)
        @show CUDA.minimum(lam)
        @show CUDA.sum(lam)
        @show CUDA.maximum(LFc)
        @show CUDA.minimum(LFc)
        @show CUDA.sum(LFc)
        @show CUDA.maximum(flux)
        @show CUDA.minimum(flux)
        @show CUDA.sum(flux)
        @show CUDA.maximum(gradf)
        @show CUDA.minimum(gradf)
        @show CUDA.sum(gradf)
        @show CUDA.maximum(rhsQ)
        @show CUDA.minimum(rhsQ)
        @show CUDA.sum(rhsQ)
        @show rhstest
        @show typeof(Q)
        @show typeof(VU)
        @show typeof(Qh)
        @show typeof(QM)
        @show typeof(QP)
        @show typeof(flux)
        @show typeof(gradf)
        @show typeof(rhsQ)
    end
    return rhsQ,rhstest
end

xq,yq,zq = (x->reshape(x,Nq_P,Np_F,K)).((xq,yq,zq))
ρ_exact(x,y,z,t) = @. 1.0f0+0.2f0*sin(pi*(x+y+z-3/2*t))
#=
# ================= #
# === Test Case === #
# ================= #
ρ = @. 1.0f0+0.2f0*sin(pi*(xq+yq+zq))
u = ones(size(xq))
v = -0.5f0*ones(size(xq))
w = ones(size(xq))
p = ones(size(xq))
=#
# ==================== #
# === Brown-Minion === #
# ==================== #
epsilon = 30
delta = .05
ρ = ones(size(xq))
u = zeros(size(xq))
#=
for i = 1:Nq_P*Np_F*K
    u[i] = (yq[i] < 0) ? tanh(epsilon*(1+abs(zq[i]))*(yq[i]+0.25)) : tanh(epsilon*(1+abs(zq[i]))*(0.25-yq[i]))
end
=#
for i = 1:Nq_P*Np_F*K
    u[i] = (yq[i] < 0) ? (1+.5*cos(pi*xq[i])*cos(pi*zq[i]))*tanh(epsilon*(yq[i]+0.25)) : (1+.5*cos(pi*xq[i])*cos(pi*zq[i]))*tanh(epsilon*(0.25-yq[i]))
end
v = @. delta*cos(2*pi*xq)
w = @. delta*cos(2*pi*xq)
Ma = 0.7
p = (1/(Ma^2*gamma))*ones(size(xq))
open("xq.txt","w") do io
    writedlm(io,Array(xq))
end
open("yq.txt","w") do io
    writedlm(io,Array(yq))
end
Q_exact(x,y,z,t) = (ρ_exact(x,y,z,t),ones(size(x)),-0.5f0*ones(size(x)),ones(size(x)),ones(size(x)))

Q = primitive_to_conservative(ρ,u,v,w,p)
Q = collect(Q)
Q_vec = zeros(Nq_P,Np_F,Nd,K)
Q_ex_vec = zeros(Nq_P,Np_F,Nd,K)
# TODO: clean up
rq2,sq2,wq2 = quad_nodes_2D(N_P+2)
Vq2 = vandermonde_2D(N_P,rq2,sq2)/VDM
xq2,yq2,zq2 = (x->Vq2*reshape(x,Np_P,Np_F*K)).((x,y,z))
Q_ex = Q_exact(xq2,yq2,zq2,T)

for k = 1:K
    for d = 1:5
        @. Q_vec[:,:,d,k] = Q[d][:,:,k]
    end
end
Q = Q_vec[:]
Q = convert.(NUM_TYPE,Q)

Q = CuArray(Q)
resQ = CUDA.fill(0.0f0,Nq*Nd*K)

################################
######   Time stepping   #######
################################
#=
# ===================== #
# ======= RK45 ======== #
# ===================== #
@time begin

@inbounds begin
for i = 1:Nsteps
    rhstest = 0.0

    for INTRK = 1:5
        if enable_test
            @show "==============="
            @show INTRK
            @show "==============="
        end
        compute_rhstest = INTRK==5
        rhsQ,rhstest = rhs(Q,ops,mesh,param,num_threads,compute_rhstest,enable_test)
        resQ .= rk4a[INTRK]*resQ+dt*rhsQ
        Q .= Q + rk4b[INTRK]*resQ
    end

    if i%10==0 || i==Nsteps
        println("Time step: $i out of $Nsteps with rhstest = $rhstest")
    end
end
end # end @inbounds

end # end @time
=#
# ========================= #
# ========= DOPRI ========= #
# ========================= #
dopri_a = zeros(NUM_TYPE,7,7)
dopri_E = zeros(NUM_TYPE,7)
dopri_a[2,1] = 1.0/5.0
dopri_a[3,1] = 3.0/40.0
dopri_a[3,2] = 9.0/40.0
dopri_a[4,1] = 44.0/45.0
dopri_a[4,2] = -56.0/15.0
dopri_a[4,3] = 32.0/9.0
dopri_a[5,1] = 19372.0/6561.0
dopri_a[5,2] = -25360.0/2187.0
dopri_a[5,3] = 64448.0/6561.0
dopri_a[5,4] = -212.0/729.0
dopri_a[6,1] = 9017.0/3168.0
dopri_a[6,2] = -355.0/33.0
dopri_a[6,3] = 46732.0/5247.0
dopri_a[6,4] = 49.0/176.0
dopri_a[6,5] = -5103.0/18656.0
dopri_a[7,1] = 35.0/384.0
dopri_a[7,2] = 0.0
dopri_a[7,3] = 500.0/1113.0
dopri_a[7,4] = 125.0/192.0
dopri_a[7,5] = -2187.0/6784.0
dopri_a[7,6] = 11.0/84.0
dopri_E[1] = dopri_a[7,1] - 5179.0/57600.0
dopri_E[2] = dopri_a[7,2]
dopri_E[3] = dopri_a[7,3] - 7571.0/16695.0
dopri_E[4] = dopri_a[7,4] - 393.0/640.0
dopri_E[5] = dopri_a[7,5] - -92097.0/339200.0
dopri_E[6] = dopri_a[7,6] - 187.0/2100.0
dopri_E[7] = dopri_a[7,7] - 1.0/40.0

Qtmp = CUDA.similar(Q)
rhsQrk = ntuple(x->CUDA.zero(Q),length(dopri_E))

errEst = 0.0
prevErrEst = 0.0

t = 0.0
i = 0
interval = 10

# false at 0.5, 1.0, 1.5, 2.0
plot_hist = [false,false,false,false]

@time begin
rhsQ,_ = rhs(Q,ops,mesh,param,num_threads,false,enable_test)
rhsQrk[1] .= rhsQ
dt0 = dt

@inbounds begin
while t < T
    rhstest = 0.0
    for INTRK = 2:7
        compute_rhstest = INTRK==7
        k = CUDA.zero(Qtmp)
        for s = 1:INTRK-1
            @. k = k+dopri_a[INTRK,s]*rhsQrk[s]
        end
        @. Qtmp = Q + dt*k
        rhsQ,rhstest = rhs(Qtmp,ops,mesh,param,num_threads,compute_rhstest,enable_test)
        @. rhsQrk[INTRK] = rhsQ
    end
    errEstVec = CUDA.zero(Qtmp)
    for s = 1:7
        @. errEstVec = errEstVec + dopri_E[s]*rhsQrk[s]
    end

    errEstScale = @. CUDA.abs(errEstVec)/(errTol*(1.0f0+CUDA.abs(Q)))
    errEst = CUDA.sum(errEstScale.*errEstScale)
    errEst = sqrt(errEst/CUDA.length(Q))
    if errEst < 1.0
        Q .= Qtmp
        global t += dt
        rhsQrk[1] .= rhsQrk[7]
    end
    order = 5
    dtnew = .8*dt*(.9/errEst)^(.4/(order+1))
    if i > 0
        dtnew *= (prevErrEst/max(1e-14,errEst))^(.3/(order+1))
    end
    global dt = max(min(10*dt0,dtnew),1e-9)
    global prevErrEst = errEst
    global i += 1
    if i % interval == 0
        println("Time step = $i, current time = $t, dt = $dtnew, errEst = $errEst, rhstest = $rhstest")
    end

    if t > 0.5 && plot_hist[1] == false
        plot_hist[1] = true
        open("array_1.txt","w") do io
            writedlm(io,Array(Q))
        end
    end

    if t > 1.0 && plot_hist[2] == false
        plot_hist[2] = true
        open("array_2.txt","w") do io
            writedlm(io,Array(Q))
        end
    end

    if t > 1.5 && plot_hist[3] == false
        plot_hist[3] = true
        open("array_3.txt","w") do io
            writedlm(io,Array(Q))
        end
    end

    if t > 2.0 && plot_hist[4] == false
        plot_hist[4] = true
        open("array_4.txt","w") do io
            writedlm(io,Array(Q))
        end
    end

    if t > T
        open("array_end.txt","w") do io
            writedlm(io,Array(Q))
        end
    end
end # end while
end
#=
open("array_end.txt","w") do io
    writedlm(io,Array(Q))
end
=#
end

#=
# Calculate L2 error
Q = Array(Q)
Q = reshape(Q,Nq_P,Np_F,Nd,K)
Q_mat = [zeros(Nq_P,Np_F,K),zeros(Nq_P,Np_F,K),zeros(Nq_P,Np_F,K),zeros(Nq_P,Np_F,K),zeros(Nq_P,Np_F,K)]
for k = 1:K
    for d = 1:5
        Q_mat[d][:,:,k] = Q[:,:,d,k]
    end
end
Pq = Array(Pq)
Q = (x->Vq2*Pq*reshape(x,Nq_P,Np_F*K)).(Q_mat)
p = pfun(Q[1],(Q[2],Q[3],Q[4]),Q[5])
Q = (Q[1],Q[2]./Q[1],Q[3]./Q[1],Q[4]./Q[1],p)
Q_ex = (x->reshape(x,length(rq2),Np_F*K)).(Q_ex)
L2_err = 0.0
for fld in 1:Nd
    global L2_err
    L2_err += sum(h*J*wq2.*(Q[fld]-Q_ex[fld]).^2)
end
println("L2err at final time T = $T is $L2_err\n")
=#
xy_val = Nd*Nh_P*Nk_tri_max
z_val = Nd*Np_F*Nk_fourier_max
println("Number of threads: $num_threads")
println("Elements in shared memory, xy flux differencing: $xy_val")
println("Elements in shared memory, z flux differencing: $z_val")

max_stars_repo_name: yiminllin/ESDG-Fourier

max_issues_repo_name: yiminllin/ESDG-Fourier

max_forks_repo_name: yiminllin/ESDG-Fourier


text:
info("Test read bam file")
data = HTSLIB.readlines("data/100.sam")
#@show data
info("Read bam successfully")

#=
info("Test write sam file")

fw = HTSLIB.open("data/test_write.sam","w","sam")
fr = HTSLIB.open("data/100.sam","r","sam")

header = unsafe_load(fr.phdr)
newheader = deepcopy(header)

phdr = convert(Ptr{HTSLIB.Header}, pointer_from_objref(newheader))

#ss = HTSLIB.strptr(fw.phdr)

fw.phdr = phdr

HTSLIB.writelines(fw,data)

HTSLIB.close(fw)
HTSLIB.close(fr)

info("writelines(fw,data) success")
=#
fw = HTSLIB.open("data/test_write2.bam","w","bam")
fr = HTSLIB.open("data/100.sam","r","sam")

header = unsafe_load(fr.phdr)
newheader = deepcopy(header)

phdr = convert(Ptr{HTSLIB.Header}, pointer_from_objref(newheader))

#ss = HTSLIB.strptr(fw.phdr)

fw.phdr = phdr

HTSLIB.writelines(fw,data)

HTSLIB.close(fw)
HTSLIB.close(fr)

info("writelines success")

ext: jl


text:
struct EM{TGV, TF<:AbstractFloat} <: FirstOrderNonMetropolisSampler
    ∇V!::TGV
    β::TF
    Δt::TF
    σ::TF
end

"""
    EM(∇V!, β, γ, M, Δt)

Set up the EM integrator for overdamped Langevin.

### Fields

* ∇V!   - In place gradient of the potential
* β     - Inverse temperature
* Δt    - Time step
"""
function EM(∇V!::TGV, β::TF, Δt::TF) where{TGV, TF<:AbstractFloat}
    σ = sqrt(2 * Δt /β);
    return EM(∇V!, β, Δt, σ)
end

mutable struct EMState{Tx} <:FirstOrderNonMetropolisSamplerState
    x::Tx
    ∇V::Tx
end

function InitState!(x₀, sampler::EM)

    ∇V = similar(x₀);
    sampler.∇V!(∇V, x₀);
    return EMState(x₀, copy(∇V));
end

function InitState(x₀, sampler::EM)

    ∇V = similar(x₀);
    sampler.∇V!(∇V, x₀);
    return EMState(deepcopy(x₀), copy(∇V));
end

function UpdateState!(state::EMState, sampler::EM)

    @. state.x = state.x - sampler.Δt * state.∇V + sampler.σ * randn();
    sampler.∇V!(state.∇V, state.x);

    state
end

max_stars_repo_name: gideonsimpson/JuBasicMD


text:
using MomentClosure
using MomentClosure: define_μ, define_M
using Test
using Catalyst

@parameters t, c₁, c₂, c₃, c₄, Ω
@variables X(t), Y(t)

S_mat = [ 1 -1  1 -1;
         -1  1  0  0]
# NOTE: only holds if combinatoric_ratelaw=false
a = [c₁*X*Y*(X-1)/Ω^2, c₂*X, c₃*Ω, c₄*X]

# checking if ReactionSystem and ReactionSystemMod are consistent

rn1 = @reaction_network begin
    (c₁/Ω^2), 2X + Y → 3X
    (c₂), X → Y
    (Ω*c₃, c₄), 0 ↔ X
end c₁ c₂ c₃ c₄ Ω

rn2 = ReactionSystemMod(t, [X, Y], [c₁, c₂, c₃, c₄, Ω], a, S_mat)

@test isequal(species(rn1), species(rn2))
@test isequal(params(rn1), params(rn2))
@test isequal(speciesmap(rn1), speciesmap(rn2))
@test isequal(paramsmap(rn1), paramsmap(rn2))
@test numspecies(rn1) == numspecies(rn2) == 2
@test numreactions(rn1) == numreactions(rn2) == 4
@test numparams(rn1) == numparams(rn2) == 5
@test netstoichmat(rn1) == netstoichmat(rn2) == S_mat
@test isequal(propensities(rn1, combinatoric_ratelaw=false), propensities(rn2))


text:
"""
---
Demodulate a BF-OFDM signal and returns the obtained T/F matrix composed of QAM symbols (without equalisation).
Based on bfofdmSigGen.jl
# --- Syntax
qamRx	= bfofdmSigDecode(sigRx,nFBMC,nOFDM,K,GI,δ,allocatedSubcarriers;posWindow="end")
# --- Input parameters
- sigRx		: Complex baseband signal to decode [Array{Complex{Float64}}]
- nFBMC		: PPN size [Int]
- nOFDM		: FFT precoder size [Int]
- K			: Overlapping factor [Int]
- GI		: CP size of precoder [Int]
- δ			: Compression rate  [Float64]
- allocatedSubcarriers : Vector of allocated subcarriers [Array{Int}]
- posWindow	  : Receiver window position. By default window is at same place as OFDM (i.e drop CP). For BF-OFDM, middle window position can also be considered with small delay spread channel (reduce ISI incuded by PPN). In that case, a phase rotation must be applied (see [1])
# --- Output parameters
- qamRx		: Decoded constellation [Array{Complex{Float64}}]
# References
- [1]	: Demmer, D.; Zakaria, R.; Gerzaguet, R.; Doré, J. & Le Ruyet, D. Study of OFDM Precoded Filter-Bank Waveforms, IEEE Transactions on Wireless Communications, 2019.
# ---
# v 1.0 - Robin Gerzaguet.
"""
function bfofdmSigDecode(sigRx,nFBMC,nOFDM,K,GI,δ,allocatedSubcarriers;posWindow="end")
	# ----------------------------------------------------
	# --- Getting number of decoded symbol and create output matrix
	# ----------------------------------------------------
	# --- PPN output symbols
	nbSymb	  = Int( length(sigRx)/nFBMC/δ +1 -K/δ );
	# --- QAM output symbols
	sizeSymb		= nOFDM + GI;
	nbSymbQamDecode = Int( nbSymb ÷ sizeSymb);
	qamRx			= zeros(Complex{Float64},length(allocatedSubcarriers),nbSymbQamDecode);
	# --- Decoding parameters
	sizeFFT			= Int( nFBMC * nOFDM * δ );
	sizeL			= Int( nFBMC * (nOFDM+GI) * δ );
	# --- Window selector
	# Defining CP as pure CP (and not half CP+CS)
	gil				= GI;				  # --- Left part appended (this is a CP)
	gir				= 0;				  # --- Right part appended (this is a CS -> 0)
	perf_GI			= Int(GI*nFBMC*δ);	  # --- Offset for Rx FFT window
	# Calculating Rx position
	giLeft		    = Int(floor(-(sizeFFT-((nOFDM+gil+gir-1)*nFBMC*δ+K*nFBMC))*δ));

	# ----------------------------------------------------
	# --- Rw Windowing
	# ----------------------------------------------------
	#TODO To be continued
	ww				  = 0;
	windowBFOFDM	  = 0;
	tmpVect			  = zeros(Complex{Float64}, sizeFFT + ww);

	# ----------------------------------------------------
	# --- Decoding
	# ----------------------------------------------------
	if posWindow == "middle"
		# ----------------------------------------------------
		# --- Best position for Rx window
		# ----------------------------------------------------
		# PPN acts as a channel FIR. But, high power taps is not the first h = . : | : .
		# As a consequence, ISI occurs (see [1] for mathematical explanation)
		# Contrary to OFDM where FFT window is get at the end of CP, for us as PPN filter raises and fall off
		# It is better to get the middle position for the Rx FFT.
		for n = 1 : 1 : nbSymbQamDecode
			# --- Getting current signal
			currSig = sigRx[ (n-1)*sizeL + giLeft - ww÷2 .+ collect((1:sizeFFT+ww)) ];
			# --- Windowing
			if windowBFOFDM == 1
				# --- Apply Windowing
				currSig	  = currSig .* filterRx;
				# --- Apply tail biting
				tmpVect[1:winLenght]		  = currSig[ end+1-winLenght:end ];
				tmpVect[ end+1-winLenght:end ]  = currSig[ 1:winLenght ];
				currSig						  = currSig + tmpVect;
				currSig						  = currSig[ 1+winLenght÷2:end-winLenght÷2 ];
			end
			# --- FFT
			# Adding a rotor term as the FFT window is not set at the end of CP but in the middle of the extended symbol
			xFFT	= fft(currSig,1) .* exp.(2im*pi.*collect(0:length(currSig)-1) .* (perf_GI-giLeft ) ./length(currSig));
			# --- OQAM decoding
			qamRx[:,n]	  = xFFT[allocatedSubcarriers];
		end
	else
		# ----------------------------------------------------
		# --- Position at the end of CP
		# ----------------------------------------------------
		# Not a clever idea for interference managment but it works (let's be honest, it is really good except in strong multipath channels).
		# Pure CP-OFDM receiver here.
		for n = 1 : 1 : nbSymbQamDecode
			# --- Getting current signal
			currSig = sigRx[ (n-1)*sizeL+perf_GI.+ collect((1:sizeFFT)) ];
			# --- FFT
			# Note that no need to rotor the signal as window is set on the pure symbol.
			xFFT    = fft(currSig,1);
			# --- Restrict to allocated subcarriers
			qamRx[:,n]	  = xFFT[allocatedSubcarriers];
		end
	end
	# --- Output on demodulated constellation
	return qamRx;
end

# ----------------------------------------------------
# --- Multiple dispatch with structure
# ----------------------------------------------------
function bfofdmSigDecode(sigId::Array{Complex{Float64}},bfofdm::StrucBFOFDM;posWindow="end")
	  return bfofdmSigDecode(sigId,bfofdm.nFBMC,bfofdm.nOFDM,bfofdm.K,bfofdm.GI,bfofdm.δ,bfofdm.allocatedSubcarriers,posWindow=posWindow);
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct ProximityPlacementGroupProperties <: SwaggerModel
    proximityPlacementGroupType::Any # spec type: Union{ Nothing, String } # spec name: proximityPlacementGroupType
    virtualMachines::Any # spec type: Union{ Nothing, Vector{SubResourceWithColocationStatus} } # spec name: virtualMachines
    virtualMachineScaleSets::Any # spec type: Union{ Nothing, Vector{SubResourceWithColocationStatus} } # spec name: virtualMachineScaleSets
    availabilitySets::Any # spec type: Union{ Nothing, Vector{SubResourceWithColocationStatus} } # spec name: availabilitySets
    colocationStatus::Any # spec type: Union{ Nothing, InstanceViewStatus } # spec name: colocationStatus

    function ProximityPlacementGroupProperties(;proximityPlacementGroupType=nothing, virtualMachines=nothing, virtualMachineScaleSets=nothing, availabilitySets=nothing, colocationStatus=nothing)
        o = new()
        validate_property(ProximityPlacementGroupProperties, Symbol("proximityPlacementGroupType"), proximityPlacementGroupType)
        setfield!(o, Symbol("proximityPlacementGroupType"), proximityPlacementGroupType)
        validate_property(ProximityPlacementGroupProperties, Symbol("virtualMachines"), virtualMachines)
        setfield!(o, Symbol("virtualMachines"), virtualMachines)
        validate_property(ProximityPlacementGroupProperties, Symbol("virtualMachineScaleSets"), virtualMachineScaleSets)
        setfield!(o, Symbol("virtualMachineScaleSets"), virtualMachineScaleSets)
        validate_property(ProximityPlacementGroupProperties, Symbol("availabilitySets"), availabilitySets)
        setfield!(o, Symbol("availabilitySets"), availabilitySets)
        validate_property(ProximityPlacementGroupProperties, Symbol("colocationStatus"), colocationStatus)
        setfield!(o, Symbol("colocationStatus"), colocationStatus)
        o
    end
end # type ProximityPlacementGroupProperties

const _property_map_ProximityPlacementGroupProperties = Dict{Symbol,Symbol}(Symbol("proximityPlacementGroupType")=>Symbol("proximityPlacementGroupType"), Symbol("virtualMachines")=>Symbol("virtualMachines"), Symbol("virtualMachineScaleSets")=>Symbol("virtualMachineScaleSets"), Symbol("availabilitySets")=>Symbol("availabilitySets"), Symbol("colocationStatus")=>Symbol("colocationStatus"))
const _property_types_ProximityPlacementGroupProperties = Dict{Symbol,String}(Symbol("proximityPlacementGroupType")=>"String", Symbol("virtualMachines")=>"Vector{SubResourceWithColocationStatus}", Symbol("virtualMachineScaleSets")=>"Vector{SubResourceWithColocationStatus}", Symbol("availabilitySets")=>"Vector{SubResourceWithColocationStatus}", Symbol("colocationStatus")=>"InstanceViewStatus")
Base.propertynames(::Type{ ProximityPlacementGroupProperties }) = collect(keys(_property_map_ProximityPlacementGroupProperties))
Swagger.property_type(::Type{ ProximityPlacementGroupProperties }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_ProximityPlacementGroupProperties[name]))}
Swagger.field_name(::Type{ ProximityPlacementGroupProperties }, property_name::Symbol) =  _property_map_ProximityPlacementGroupProperties[property_name]

const _allowed_ProximityPlacementGroupProperties_proximityPlacementGroupType = ["Standard", "Ultra"]

function check_required(o::ProximityPlacementGroupProperties)
    true
end

function validate_property(::Type{ ProximityPlacementGroupProperties }, name::Symbol, val)
    if name === Symbol("proximityPlacementGroupType")
        Swagger.validate_param(name, "ProximityPlacementGroupProperties", :enum, val, _allowed_ProximityPlacementGroupProperties_proximityPlacementGroupType)
    end
end


text:
module Elemental

using Distributed
using DistributedArrays
using LinearAlgebra

import Base: *, \
import Base: convert, copy, copy!, eltype, fill!, getindex, hcat, inv, length, pointer, print, setindex!, similar, size

using Elemental_jll: libEl

include("error.jl")

function Init()
    argc = Cint[0]
    argv = Vector{String}[String[""]]
    ElError(ccall((:ElInitialize, libEl), Cint,
                (Ptr{Cint}, Ptr{Ptr{Ptr{UInt8}}}), argc, pointer(argv)))
    return nothing
end

function Initialized()
    active = Ref(zero(Cint))
    ElError(ccall((:ElInitialized, libEl), Cuint, (Ref{Cint},), active))
    return active[] == 1
end

function Finalize()
    ElError(ccall((:ElFinalize, libEl), Cint, ()))
    return nothing
end

function __init__()
    # ccall(:jl_, Cvoid, (Any,), "starting up!")
    Init()
    DefaultGrid[] = Grid()
    atexit() do
        # ccall(:jl_, Cvoid, (Any,), "closing down!")
        Initialized() && Finalize()
    end
end

include("core/types.jl")
include("mpi.jl")
include("core/matrix.jl")
include("core/grid.jl")
include("core/sparsematrix.jl")
include("core/distmatrix.jl")
include("core/distsparsematrix.jl")
include("core/distmultivec.jl")
include("matrices.jl")
include("io.jl")
include("blas_like/level1.jl")
include("blas_like/level2.jl")
include("blas_like/level3.jl")
include("lapack_like/euclidean_min.jl")
include("lapack_like/condense.jl")
include("lapack_like/factor.jl")
include("lapack_like/props.jl")
include("lapack_like/funcs.jl")
include("lapack_like/solve.jl")
include("lapack_like/spectral.jl")
include("optimization/solvers.jl")
include("optimization/models.jl")
include("julia/generic.jl")
include("julia/darray.jl")

end # module


text:
print("RUNNING TESTS:    ")
include("test.jl")
println("[OK]")


text:
function gradexpr(ex::SymbolicVariable, wrt::SymbolicVariable)
    return ex == wrt ? 1 : 0
end

gradexpr(ex::Number, wrt::SymbolicVariable) = 0

function gradexpr(ex::Expr, wrt::Any)
    if ex.head != :call
        error("Unrecognized expression $ex")
    end
    simplify(
        gradexpr(
            SymbolParameter(ex.args[1]),
            ex.args[2:end],
            wrt
        )
    )
end

function gradexpr{T}(x::SymbolParameter{T}, args::Any, wrt::Any)
    error(
        @sprintf(
            "Derivative of function %s not supported",
            string(T)
        )
    )
end

# The Power Rule
function gradexpr(::SymbolParameter{:^}, args::Any, wrt::Any)
    x, y = args[1], args[2]
    xp, yp = gradexpr(x, wrt), gradexpr(y, wrt)
    if xp == 0 && yp == 0
        return 0
    elseif yp == 0
        return :($y * $xp * ($x^($y - 1)))
    else
        return :($x^$y * ($xp * $y / $x + $yp * log($x)))
    end
end

function gradexpr(::SymbolParameter{:+}, args::Any, wrt::Any)
    termdiffs = {:+}
    for y in args
        x = gradexpr(y, wrt)
        if x != 0
            push!(termdiffs, x)
        end
    end
    n = length(termdiffs)
    if n == 1
        return 0
    elseif n == 2
        return termdiffs[2]
    else
        return Expr(:call, termdiffs...)
    end
end

function gradexpr(::SymbolParameter{:-}, args::Any, wrt::Any)
    termdiffs = {:-}
    # first term is special, can't be dropped
    term1 = gradexpr(args[1], wrt)
    push!(termdiffs, term1)
    for y in args[2:end]
        x = gradexpr(y, wrt)
        if x != 0
            push!(termdiffs, x)
        end
    end
    n = length(termdiffs)
    if term1 != 0 && n == 2 && length(args) >= 2
        # if all of the terms but the first disappeared, we return the first
        return term1
    elseif term1 == 0 && n == 2
        return 0
    else
        return Expr(:call, termdiffs...)
    end
end

# The Product Rule
# d/dx (f * g) = (d/dx f) * g + f * (d/dx g)
# d/dx (f * g * h) = (d/dx f) * g * h + f * (d/dx g) * h + ...
function gradexpr(::SymbolParameter{:*}, args::Any, wrt::Any)
    n = length(args)
    resargs = Array(Any, n)
    for i in 1:n
       newargs = Array(Any, n)
       for j in 1:n
           if j == i
               newargs[j] = gradexpr(args[j], wrt)
           else
               newargs[j] = args[j]
           end
       end
       resargs[i] = Expr(:call, :*, newargs...)
    end
    return Expr(:call, :+, resargs...)
end

# The Quotient Rule
# d/dx (f / g) = ((d/dx f) * g - f * (d/dx g)) / g^2
function gradexpr(::SymbolParameter{:/}, args::Any, wrt::Any)
    x, y = args[1], args[2]
    xp, yp = gradexpr(x, wrt), gradexpr(y, wrt)
    if xp == 0 && yp == 0
        return 0
    elseif xp == 0
        return :(-$yp * $x / $y^2)
    elseif yp == 0
        return :($xp / $y)
    else
        return :(($xp * $y - $x * $yp) / $y^2)
    end
end

# This table is used in other packages, and if someone changes it they should
# notify the following packages:
# * https://github.com/scidom/DualNumbers.jl
derivative_rules = [
    ( :sqrt,        :(  xp / 2 / sqrt(x)                         ))
    ( :cbrt,        :(  xp / 3 / cbrt(x)^2                       ))
    ( :square,      :(  xp * 2 * x                               ))
    ( :log,         :(  xp / x                                   ))
    ( :log10,       :(  xp / x / log(10)                         ))
    ( :log2,        :(  xp / x / log(2)                          ))
    ( :log1p,       :(  xp / (x + 1)                             ))
    ( :exp,         :(  xp * exp(x)                              ))
    ( :exp2,        :(  xp * log(2) * exp2(x)                    ))
    ( :expm1,       :(  xp * exp(x)                              ))
    ( :sin,         :(  xp * cos(x)                              ))
    ( :cos,         :( -xp * sin(x)                              ))
    ( :tan,         :(  xp * (1 + tan(x)^2)                      ))
    ( :sec,         :(  xp * sec(x) * tan(x)                     ))
    ( :csc,         :( -xp * csc(x) * cot(x)                     ))
    ( :cot,         :( -xp * (1 + cot(x)^2)                      ))
    ( :sind,        :(  xp * pi / 180 * cosd(x)                  ))
    ( :cosd,        :( -xp * pi / 180 * sind(x)                  ))
    ( :tand,        :(  xp * pi / 180 * (1 + tand(x)^2)          ))
    ( :secd,        :(  xp * pi / 180 * secd(x) * tand(x)        ))
    ( :cscd,        :( -xp * pi / 180 * cscd(x) * cotd(x)        ))
    ( :cotd,        :( -xp * pi / 180 * (1 + cotd(x)^2)          ))
    ( :asin,        :(  xp / sqrt(1 - x^2)                       ))
    ( :acos,        :( -xp / sqrt(1 - x^2)                       ))
    ( :atan,        :(  xp / (1 + x^2)                           ))
    ( :asec,        :(  xp / abs(x) / sqrt(x^2 - 1)              ))
    ( :acsc,        :( -xp / abs(x) / sqrt(x^2 - 1)              ))
    ( :acot,        :( -xp / (1 + x^2)                           ))
    ( :asind,       :(  xp * 180 / pi / sqrt(1 - x^2)            ))
    ( :acosd,       :( -xp * 180 / pi / sqrt(1 - x^2)            ))
    ( :atand,       :(  xp * 180 / pi / (1 + x^2)                ))
    ( :asecd,       :(  xp * 180 / pi / abs(x) / sqrt(x^2 - 1)   ))
    ( :acscd,       :( -xp * 180 / pi / abs(x) / sqrt(x^2 - 1)   ))
    ( :acotd,       :( -xp * 180 / pi / (1 + x^2)                ))
    ( :sinh,        :(  xp * cosh(x)                             ))
    ( :cosh,        :(  xp * sinh(x)                             ))
    ( :tanh,        :(  xp * sech(x)^2                           ))
    ( :sech,        :( -xp * tanh(x) * sech(x)                   ))
    ( :csch,        :( -xp * coth(x) * csch(x)                   ))
    ( :coth,        :( -xp * csch(x)^2                           ))
    ( :asinh,       :(  xp / sqrt(x^2 + 1)                       ))
    ( :acosh,       :(  xp / sqrt(x^2 - 1)                       ))
    ( :atanh,       :(  xp / (1 - x^2)                           ))
    ( :asech,       :( -xp / x / sqrt(1 - x^2)                   ))
    ( :acsch,       :( -xp / abs(x) / sqrt(1 + x^2)              ))
    ( :acoth,       :(  xp / (1 - x^2)                           ))
    ( :erf,         :(  xp * 2 * exp(-square(x)) / sqrt(pi)      ))
    ( :erfc,        :( -xp * 2 * exp(-square(x)) / sqrt(pi)      ))
    ( :erfi,        :(  xp * 2 * exp(square(x)) / sqrt(pi)       ))
    ( :gamma,       :(  xp * digamma(x) * gamma(x)               ))
    ( :lgamma,      :(  xp * digamma(x)                          ))
    # note: only covers the 1-arg version of :airy
    ( :airy,        :(  xp * airyprime(x)                        ))
    ( :airyprime,   :(  xp * airy(2, x)                          ))
    ( :airyai,      :(  xp * airyaiprime(x)                      ))
    ( :airybi,      :(  xp * airybiprime(x)                      ))
    ( :airyaiprime, :(  xp * x * airyai(x)                       ))
    ( :airybiprime, :(  xp * x * airybi(x)                       ))
    ( :besselj0,    :( -xp * besselj1(x)                         ))
    ( :besselj1,    :(  xp * (besselj0(x) - besselj(2, x)) / 2   ))
    ( :bessely0,    :( -xp * bessely1(x)                         ))
    ( :bessely1,    :(  xp * (bessely0(x) - bessely(2, x)) / 2   ))
    ## ( :erfcx,   :(  xp * (2 * x * erfcx(x) - 2 / sqrt(pi))   ))  # uncertain
    ## ( :dawson,  :(  xp * (1 - 2x * dawson(x))                ))  # uncertain
]

for (funsym, exp) in derivative_rules
    @eval function gradexpr(
        ::SymbolParameter{$(Meta.quot(funsym))},
        args::Any,
        wrt::Any
    )
        x = args[1]
        xp = gradexpr(x, wrt)
        if xp != 0
            return @sexpr($exp)
        else
            return 0
        end
    end
end

derivative_rules_bessel = [
    ( :besselj,    :(  xp * (besselj(nu - 1, x) - besselj(nu + 1, x)) / 2   ))
    ( :besseli,    :(  xp * (besseli(nu - 1, x) + besseli(nu + 1, x)) / 2   ))
    ( :bessely,    :(  xp * (bessely(nu - 1, x) - bessely(nu + 1, x)) / 2   ))
    ( :besselk,    :( -xp * (besselk(nu - 1, x) + besselk(nu + 1, x)) / 2   ))
    ( :hankelh1,   :(  xp * (hankelh1(nu - 1, x) - hankelh1(nu + 1, x)) / 2 ))
    ( :hankelh2,   :(  xp * (hankelh2(nu - 1, x) - hankelh2(nu + 1, x)) / 2 ))
]

# 2-argument bessel functions
for (funsym, exp) in derivative_rules_bessel
    @eval function gradexpr(
        ::SymbolParameter{$(Meta.quot(funsym))},
        args::Any,
        wrt::Any
    )
        nu, x = args[1], args[2]
        xp = gradexpr(x, wrt)
        if xp != 0
            return @sexpr($exp)
        else
            return 0
        end
    end
end

### Other functions from julia/base/math.jl we might want to define
### derivatives for. Some have two arguments.

## atan2
## hypot
## beta, lbeta, eta, zeta, digamma

function gradexpr(ex::Expr, targets::Vector{Symbol})
    n = length(targets)
    exprs = Array(Any, n)
    for i in 1:n
        exprs[i] = gradexpr(ex, targets[i])
    end
    return exprs
end


text:
using TMLEEpistasis

include(joinpath(pkgdir(TMLEEpistasis), "test", "runtests.jl"))


text:
import Core.Compiler: abstract_call_gf_by_type, abstract_call
using Core.Compiler: Const, isconstType, argtypes_to_type, tuple_tfunc, Const,
    getfield_tfunc, _methods_by_ftype, VarTable, cache_lookup, nfields_tfunc
using Base.Meta

function abstract_call_gf_by_type(interp::ADInterpreter, @nospecialize(f), argtypes::Vector{Any}, @nospecialize(atype), sv::InferenceState, max_methods::Int = InferenceParams(interp).MAX_METHODS)
    if f isa ∂⃖
        inner_argtypes = argtypes[2:end]
        ft = inner_argtypes[1]
        f = singleton_type(ft)
        rinterp = raise_level(interp)
        call = abstract_call_gf_by_type(rinterp, f, inner_argtypes, argtypes_to_type(inner_argtypes), sv, max_methods)
        if isa(call.info, MethodMatchInfo)
            if length(call.info.results.matches) == 0
                @show inner_argtypes
                error()
            end
            mi = specialize_method(call.info.results.matches[1], true)
            ci = get(code_cache(rinterp), mi, nothing)
            clos = AbstractCompClosure(rinterp.current_level, 1, call.info, ci.inferred.stmts.info)
        elseif isa(call.info, RRuleInfo)
            if rinterp.current_level == 1
                clos = getfield_tfunc(call.info.rrule_rt, Const(2))
            else
                name = call.info.info.results.matches[1].method.sig.parameters[2].name.mt.name
                clos = PrimClosure(name, rinterp.current_level - 1, 1, getfield_tfunc(call.info.rrule_rt, Const(2)), call.info, nothing)
            end
        end
        rt2 = tuple_tfunc(Any[call.rt, clos])
        return CallMeta(rt2, RecurseInfo(call.info))
    end

    # Check if there is a rrule for this function
    if interp.current_level != 0 && f !== ChainRules.rrule
        rrule_argtypes = Any[Const(ChainRules.rrule); argtypes]
        rrule_atype = argtypes_to_type(rrule_argtypes)
        # In general we want the forward type of an rrule'd function to match
        # what the function itself would have returned, but let's support this
        # not being the case.
        if f == accum
            error()
        end
        call = abstract_call_gf_by_type(lower_level(interp), ChainRules.rrule, rrule_argtypes, rrule_atype, sv, -1)
        if call.rt != Const(nothing)
            @show (f,call.rt)
            return CallMeta(getfield_tfunc(call.rt, Const(1)), RRuleInfo(call.rt, call.info))
        end
    end

    ret = invoke(abstract_call_gf_by_type,
        Tuple{AbstractInterpreter, Any, Vector{Any}, Any, InferenceState, Int},
        interp, f, argtypes, atype, sv, max_methods)

    return ret
end

function abstract_accum(interp::AbstractInterpreter, args::Vector{Any}, sv::InferenceState)
    args = filter(x->!(widenconst(x) <: Union{ZeroTangent, NoTangent}), args)

    if length(args) == 0
        return CallMeta(ZeroTangent, nothing)
    end

    if length(args) == 1
        return CallMeta(args[1], nothing)
    end

    rtype = reduce(tmerge, args)
    if widenconst(rtype) <: Tuple
        targs = Any[]
        for i = 1:nfields_tfunc(rtype).val
            push!(targs, abstract_accum(interp, Any[getfield_tfunc(arg, Const(i)) for arg in args], sv).rt)
        end
        return CallMeta(tuple_tfunc(targs), nothing)
    end
    call = abstract_call(change_level(interp, 0), nothing, Any[typeof(accum), args...],
        sv::InferenceState)
    return call
end

function repackage_apply_rt(info, Δ, argtypes)
    argwise = Any[NoTangent, NoTangent, getfield_prop_zero_tfunc(Δ, Const(1))]
    curarg = 2
    # Repackage this according to the arguments
    for argt in argtypes
        arg_rt = widenconst(argt)
        #@show arg_rt
        if arg_rt <: Tuple
            n = length(arg_rt.parameters)
            tts = Any[]
            for j = curarg:(curarg+n-1)
                push!(tts, getfield_prop_zero_tfunc(Δ, Const(j)))
            end
            push!(argwise, tuple_tfunc(tts))
            curarg += n
        else
            error()
        end
    end
    final_rt = tuple_tfunc(argwise)
    return final_rt
end

function infer_cc_backward(interp::ADInterpreter, cc::AbstractCompClosure, @nospecialize(cc_Δ), sv::InferenceState)
    @show ("enter", cc_Δ, cc)

    mi = specialize_method(cc.primal_info.results.matches[1], true)
    ni = change_level(interp, cc.order)
    ci = get(code_cache(ni), mi, nothing)
    primal = ci.inferred

    @show primal

    function derive_closure_type(info)
        isa(info, CallMeta) && (info = info.info)
        if isa(info, RRuleInfo)
            if cc.order == 1
                return (0, getfield_tfunc(info.rrule_rt, Const(2)))
            else
                tn = info.info.results.matches[1].method.sig.parameters[2].name
                name = isdefined(tn, :mt) ? tn.mt.name : tn.name
                return (cc.order - 1, PrimClosure(name, cc.order - 1, 1, getfield_tfunc(info.rrule_rt, Const(2)), info, nothing))
            end
        elseif isa(info, MethodMatchInfo)
            mi′ = specialize_method(info.results.matches[1], true)
            ci′ = get(code_cache(ni), mi′, nothing)

            return (0, AbstractCompClosure(cc.order, cc.seq, info, ci′.inferred.stmts.info))
        elseif isa(info, CompClosInfo)
            return (info.clos.order, AbstractCompClosure(info.clos.order, info.clos.seq + 1, info.clos.primal_info, info.infos))
        elseif isa(info, PrimClosInfo)
            return (info.next.order, info.next)
        elseif isa(info, ReifyInfo)
            return derive_closure_type(info.info)
        else
            @show info
            error()
        end
    end

    ssa_accums = Vector{Union{Nothing, CallMeta}}(undef, length(primal.stmts))
    ssa_infos = Union{Nothing, CallMeta}[nothing for i = 1:length(primal.stmts)]
    arg_accums = Vector{Union{Nothing, CallMeta}}(undef, length(primal.argtypes))

    arg_typs = Any[Any[] for i = 1:length(primal.argtypes)]
    ssa_typs = Any[Any[] for i = 1:length(primal.stmts)]
    function accum!(use, val)
        if val == Union{}
            @show (use, val)
            @show cc.seq
            error()
        end
        if isa(use, SSAValue)
            push!(ssa_typs[use.id], val)
        elseif isa(use, Argument)
            push!(arg_typs[use.n], val)
        end
    end

    function bail!(inst)
        @show ("bail", inst)
        error()
        for i = 1:length(inst.args)
            accum!(inst.args[i], Any)
        end
    end

    function arg_accum!(inst, rt)
        if rt === Union{}
            @show primal
            @show inst
            @show rt
            error()
        end
        for i = 1:length(inst.args)
            accum!(inst.args[i], getfield_tfunc(rt, Const(i)))
        end
    end

    for i = length(primal.stmts):-1:1
        inst = primal.stmts[i][:inst]

        if isa(inst, ReturnNode)
            accum!(inst.val, cc_Δ)
            continue
        elseif isa(inst, GlobalRef) || isexpr(inst, :static_parameter)
            continue
        else
            if !(isexpr(inst, :call) || isexpr(inst, :new))
                @show mi
                @show inst
                @show primal
                error()
            end
        end

        accum_call = abstract_accum(change_level(interp, 0), ssa_typs[i], sv)
        ssa_accums[i] = accum_call
        Δ = accum_call.rt
        if Δ == Union{}
            error()
        end

        if isexpr(inst, :new)
            for i = 2:length(inst.args)
                # TODO: This is wrong
                rt = getfield_tfunc(Δ, Const(i-1))
                # Struct gradients are allowed to be sparse
                rt === Union{} && (rt = ZeroTangent)
                accum!(inst.args[i], rt)
            end
            accum!(inst.args[1], ChainRules.NoTangent)
            continue
        end

        info = cc.prev_seq_infos[i]
        isa(info, CallMeta) && (info = info.info)

        call_info = info
        while isa(call_info, UnionSplitApplyCallInfo)
            @assert length(info.infos) == 1
            call_info = call_info.infos[1].call
        end

        if call_info === nothing
            ft = argextype(inst.args[1], primal, primal.sptypes)
            f = singleton_type(ft)
            if isa(f, Core.Builtin)
                call = CallMeta(backwards_tfunc(f, primal, inst, Δ), nothing)
            else
                bail!(inst)
                continue
            end
        else
            if cc.seq == 1 && isa(call_info, CompClosInfo)
                call_info = ReifyInfo(call_info)
            end

            if isa(call_info, RecurseInfo)
                clos = getfield_tfunc(Δ, Const(2))
                arg = getfield_tfunc(Δ, Const(1))
                call = abstract_call(interp, nothing, Any[clos, arg], sv)
                # No derivative wrt the functor
                call = CallMeta(tuple_tfunc(Any[NoTangent; tuple_type_fields(call.rt)...]), ReifyInfo(call.info))
            else
                (level, close) = derive_closure_type(call_info)
                call = abstract_call(change_level(interp, level), nothing, Any[close, Δ], sv)
            end
        end

        if isa(info, UnionSplitApplyCallInfo)
            argts = Any[argextype(inst.args[i], primal, primal.sptypes) for i = 4:length(inst.args)]
            call = CallMeta(repackage_apply_rt(info, call.rt, argts),
                UnionSplitApplyCallInfo([ApplyCallInfo(call.info)]))
        end

        if isa(call_info, ReifyInfo)
            new_rt = tuple_tfunc(Any[derive_closure_type(call.info)[2]; call.rt])
            call = CallMeta(new_rt, RecurseInfo(call.info))
        end

        if call.rt === Union{}
            @show inst
            @show Δ
            @show cc_Δ
            @show mi
            @show primal
            @show cc.order
            @show close
            error()
        end

        arg_accum!(inst, call.rt)
        ssa_infos[i] = call
        continue
    end

    tup_elemns = Any[]
    for (i, this_arg_typs) in enumerate(arg_typs)
        let tup_push! = (mi.def.isva && i == length(arg_typs)) ? (a, t)->append!(a, tuple_type_fields(t)) : push!
            if length(this_arg_typs) <= 1
                push!(arg_accums, nothing)
                length(this_arg_typs) == 0 && tup_push!(tup_elemns, ChainRules.NoTangent)
                length(this_arg_typs) == 1 && tup_push!(tup_elemns, this_arg_typs[1])
                continue
            end
            accum_call = abstract_accum(interp, this_arg_typs, sv)
            if accum_call.rt == Union{}
                @show accum_call.rt
                return CallMeta(Union{}, false)
            end
            push!(arg_accums, accum_call)
            tup_push!(tup_elemns, accum_call.rt)
        end
    end

    rt = tuple_tfunc(Any[tup_elemns...])
    @show (cc, rt)
    return CallMeta(rt, CompClosInfo(cc, ssa_infos))
end

function infer_cc_forward(interp::ADInterpreter, cc::AbstractCompClosure, @nospecialize(cc_Δ), sv::InferenceState)
    @show ("enter", cc_Δ, cc)

    primal_mi = specialize_method(cc.primal_info.results.matches[1], true)
    ni = change_level(interp, cc.order)
    ci = get(code_cache(ni), primal_mi, nothing)
    primal = ci.inferred

    function derive_closure_type(info)
        isa(info, CallMeta) && (info = info.info)
        if isa(info, RRuleInfo)
            return (cc.order - 1, getfield_tfunc(info.rrule_rt, Const(2)))
        elseif isa(info, MethodMatchInfo)
            mi′ = specialize_method(info.results.matches[1], true)
            ci′ = get(code_cache(lower_level(ni)), mi′, nothing)

            return (0, AbstractCompClosure(cc.order - 1, 1, info, ci′.inferred.stmts.info))
        elseif isa(info, CompClosInfo)
            return (info.clos.order, AbstractCompClosure(info.clos.order, info.clos.seq + 1, info.clos.primal_info, info.infos))
        elseif isa(info, PrimClosInfo)
            return (info.next.order, info.next)
        elseif isa(info, ReifyInfo)
            return derive_closure_type(info.info)
        else
            @show info
            error()
        end
    end

    function accum_arg(arg)
        if isa(arg, SSAValue)
            if !isassigned(accums, arg.id)
                @show primal
                @show arg
                error()
            end
            return accums[arg.id]
        elseif isa(arg, Argument)
            argn = arg.n
            if primal_mi.def.isva && argn == length(primal.argtypes)
                return tuple_tfunc(Any[tuple_type_fields(cc_Δ)[argn:end]...])
            else
                return getfield_tfunc(cc_Δ, Const(argn))
            end
        else
            return NoTangent
        end
    end

    ssa_accums = Vector{Union{Nothing, CallMeta}}(undef, length(primal.stmts))
    ssa_infos = Vector{Union{Nothing, CallMeta}}(undef, length(primal.stmts))

    accums = Vector{Any}(undef, length(primal.stmts))
    vals = Vector{Union{Nothing, CallMeta}}(undef, length(primal.stmts))

    for i = 1:length(primal.stmts)
        inst = primal.stmts[i][:inst]
        info = cc.prev_seq_infos[i]

        if isa(inst, GlobalRef)
            accums[i] = NoTangent
            continue
        end

        if isa(inst, ReturnNode)
            rt = accum_arg(inst.val)
            return CallMeta(rt, CompClosInfo(cc, ssa_infos))
        end

        args = Any[]
        for i = 1:length(inst.args)
            push!(args, accum_arg(inst.args[i]))
        end

        if isexpr(inst, :new)
            T, exact = instanceof_tfunc(argextype(inst.args[1], primal, primal.sptypes))
            # Special case for empty structs
            if isempty(fieldnames(T))
                accums[i] = ZeroTangent
                continue
            end
            Δ = NamedTuple{fieldnames(T), widenconst(tuple_tfunc(args[2:end]))}
            accums[i] = Δ
            continue
        end

        Δ = tuple_tfunc(args)

        if info === nothing
            ssa_infos[i] = nothing
            ft = argextype(inst.args[1], primal, primal.sptypes)
            f = singleton_type(ft)
            if isa(f, Core.Builtin)
                accums[i] = forward_tfunc(f, primal, inst, Δ)
                continue
            else
                bail!(inst)
                continue
            end
        end

        ssa_accums[i] = nothing
        isa(info, CallMeta) && (info = info.info)

        call_info = info
        while isa(call_info, UnionSplitApplyCallInfo)
            @assert length(info.infos) == 1
            call_info = call_info.infos[1].call

            # Tranforms Δ
            args = Any[]
            push!(args, tuple_type_fields(Δ)[3])
            for arg in tuple_type_fields(Δ)[4:end]
                append!(args, tuple_type_fields(arg))
            end
            Δ = tuple_tfunc(args)
        end

        if isa(call_info, ReifyInfo)
            # TODO: Is this always the case?
            Δ = tuple_tfunc(Any[tuple_type_fields(Δ)[2:end]...])
        end

        if isa(call_info, RecurseInfo)
            clos = getfield_tfunc(Δ, Const(1))
            arg = getfield_tfunc(Δ, Const(2))
            call = abstract_call(interp, nothing, Any[clos, arg], sv)
            # No derivative wrt the functor
            call = CallMeta(tuple_tfunc(Any[NoTangent; tuple_type_fields(call.rt)...]), ReifyInfo(call.info))
            #error()
        else
            (level, clos) = derive_closure_type(call_info)

            call = abstract_call(change_level(interp, level), nothing, Any[clos, Δ], sv)
        end

        if isa(call_info, ReifyInfo)
            new_rt = tuple_tfunc(Any[call.rt; derive_closure_type(call.info)[2]])
            call = CallMeta(new_rt, RecurseInfo())
        end

        if isa(info, UnionSplitApplyCallInfo)
            call = CallMeta(call.rt, UnionSplitApplyCallInfo([ApplyCallInfo(call.info)]))
        end

        accums[i] = call.rt
        @show (i, accums[i])
        ssa_infos[i] = call
    end

    error()
end

function infer_comp_closure(interp::ADInterpreter, cc::AbstractCompClosure, @nospecialize(Δ), sv::InferenceState)
    if cc.seq & 1 != 0
        infer_cc_backward(interp, cc, Δ, sv)
    else
        infer_cc_forward(interp, cc, Δ, sv)
    end
end

function infer_prim_closure(interp::ADInterpreter, pc::PrimClosure, @nospecialize(Δ), sv::InferenceState)
    @show ("enter", pc)

    if pc.seq == 1
        call = abstract_call(change_level(interp, pc.order), nothing, Any[pc.dual, Δ], sv)
        rt = call.rt
        @show (pc, Δ, rt)
        return CallMeta(call.rt, PrimClosInfo(PrimClosure(pc.name, pc.order, pc.seq + 1, nothing, call.info, pc.info_below)))
    elseif pc.seq == 2
        ni = change_level(interp, pc.order)
        mi′ = specialize_method(pc.info_below.results.matches[1], true)
        ci′ = get(code_cache(ni), mi′, nothing)
        cc = AbstractCompClosure(pc.order, 1, pc.info_below, ci′.inferred.stmts.info)
        call = infer_comp_closure(ni, cc, Δ, sv)
        rt = getfield_tfunc(call.rt, Const(2))
        @show (pc, Δ, rt)
        return CallMeta(rt,
            PrimClosInfo(PrimClosure(pc.name, pc.order, pc.seq + 1, getfield_tfunc(call.rt, Const(1)), call.info, pc.info_carried)))
    elseif pc.seq == 3
        ni = change_level(interp, pc.order)
        mi′ = specialize_method(pc.info_carried.info.results.matches[1], true)
        ci′ = get(code_cache(ni), mi′, nothing)
        clos = AbstractCompClosure(pc.order, 1, pc.info_carried.info, ci′.inferred.stmts.info)
        call = abstract_call(change_level(interp, pc.order), nothing,
            Any[clos, tuple_tfunc(Any[Δ, pc.dual])], sv)
        rt = tuple_tfunc(Any[tuple_type_fields(call.rt)[2:end]...])
        @show (pc, Δ, rt)
        return CallMeta(rt,
            PrimClosInfo(PrimClosure(pc.name, pc.order, pc.seq + 1, nothing, call.info, pc.info_below)))
    elseif mod(pc.seq, 4) == 0
        info = pc.info_below
        clos = AbstractCompClosure(info.clos.order, info.clos.seq + 1, info.clos.primal_info, info.infos)

        # Add back gradient w.r.t. rrule
        Δ = tuple_tfunc(Any[NoTangent, tuple_type_fields(Δ)...])
        call = abstract_call(change_level(interp, pc.order), nothing, Any[clos, Δ], sv)
        rt = getfield_tfunc(call.rt, Const(1))
        @show (pc, Δ, rt)
        return CallMeta(rt, PrimClosInfo(PrimClosure(pc.name, pc.order, pc.seq + 1, getfield_tfunc(call.rt, Const(2)), call.info, pc.info_carried)))
    elseif mod(pc.seq, 4) == 1
        info = pc.info_carried
        clos = AbstractCompClosure(info.clos.order, info.clos.seq + 1, info.clos.primal_info, info.infos)
        call = abstract_call(change_level(interp, pc.order), nothing, Any[clos, tuple_tfunc(Any[pc.dual, Δ])], sv)
        rt = call.rt
        @show (pc, Δ, rt)
        return CallMeta(call.rt, PrimClosInfo(PrimClosure(pc.name, pc.order, pc.seq + 1, nothing, call.info, pc.info_below)))
    elseif mod(pc.seq, 4) == 2
        info = pc.info_below
        clos = AbstractCompClosure(info.clos.order, info.clos.seq + 1, info.clos.primal_info, info.infos)
        call = abstract_call(change_level(interp, pc.order), nothing, Any[clos, Δ], sv)
        rt = getfield_tfunc(call.rt, Const(2))
        @show (pc, Δ, rt)
        return CallMeta(rt,
            PrimClosInfo(PrimClosure(pc.name, pc.order, pc.seq + 1, getfield_tfunc(call.rt, Const(1)), call.info, pc.info_carried)))
    elseif mod(pc.seq, 4) == 3
        info = pc.info_carried
        clos = AbstractCompClosure(info.clos.order, info.clos.seq + 1, info.clos.primal_info, info.infos)
        call = abstract_call(change_level(interp, pc.order), nothing, Any[clos, tuple_tfunc(Any[Δ, pc.dual])], sv)
        rt = tuple_tfunc(Any[tuple_type_fields(call.rt)[2:end]...])
        @show (pc, Δ, rt)
        return CallMeta(rt,
            PrimClosInfo(PrimClosure(pc.name, pc.order, pc.seq + 1, nothing, call.info, pc.info_below)))
    end
    error()
end

function abstract_call(interp::ADInterpreter, fargs::Union{Nothing,Vector{Any}}, argtypes::Vector{Any},
    sv::InferenceState, max_methods::Int = InferenceParams(interp).MAX_METHODS)

    if isa(argtypes[1], AbstractCompClosure)
        if length(argtypes) !== 2
            error()
            return CallMeta(Union{}, false)
        end
        return infer_comp_closure(interp, argtypes[1], argtypes[2], sv)
    elseif isa(argtypes[1], PrimClosure)
        return infer_prim_closure(interp, argtypes[1], argtypes[2], sv)
    end

    rt = invoke(abstract_call, Tuple{AbstractInterpreter, Union{Nothing, Vector{Any}}, Vector{Any}, InferenceState, Int64},
        interp, fargs, argtypes, sv, max_methods)

    return rt
end


text:
using TaskGraphs
# NOTE that the experiments described in "Optimal Sequential Task Assignment and
# Path Finding for Multi-Agent Robotic Assembly Planning", Brown et al.
# were performed using Gurobi as the black box MILP solver. To use the default
# (GLPK), simply comment out the following two lines.
using Gurobi
set_default_milp_optimizer!(Gurobi.Optimizer)

## ICRA experiments
# You may want to redefine base_dis
# -------------------------- #
base_dir            = joinpath("/scratch/task_graphs_experiments")
# -------------------------- #
problem_dir         = joinpath(base_dir,"pctapf_problem_instances")
base_results_path   = joinpath(base_dir,"pctapf_results")

feats = [
    RunTime(),
    FeasibleFlag(),
    OptimalFlag(),
    OptimalityGap(),
    SolutionCost(),
    PrimaryCost(),
    SolutionAdjacencyMatrix(),
]
solver_configs = [
    (
        solver = NBSSolver(
            assignment_model = TaskGraphsMILPSolver(AssignmentMILP()),
            path_planner = CBSSolver(ISPS()),
            ),
        results_path = joinpath(base_results_path,"AssignmentMILP"),
        feats = feats,
        objective = MakeSpan(),
    ),
]

base_config = Dict(
    :env_id => "env_2",
    :num_trials => 16,
    :max_parents => 3,
    :depth_bias => 0.4,
    :dt_min => 0,
    :dt_max => 0,
    :dt_collect => 0,
    :dt_deliver => 0,
)
size_configs = [Dict(:M => m, :N => n) for (n,m) in Base.Iterators.product(
    [10,20,30,40],[10,20,30,40,50,60]
)][:]
problem_configs = map(d->merge(d,base_config), size_configs)

# loader = PCTA_Loader() # to test task assignment only
loader = PCTAPF_Loader()
add_env!(loader,"env_2",init_env_2())
write_problems!(loader,problem_configs,problem_dir)
for solver_config in solver_configs
    set_runtime_limit!(solver_config.solver,100)
    warmup(loader,solver_config,problem_dir)
    run_profiling(loader,solver_config,problem_dir)
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct ConnectivityParameters <: SwaggerModel
    source::Any # spec type: Union{ Nothing, ConnectivitySource } # spec name: source
    destination::Any # spec type: Union{ Nothing, ConnectivityDestination } # spec name: destination

    function ConnectivityParameters(;source=nothing, destination=nothing)
        o = new()
        validate_property(ConnectivityParameters, Symbol("source"), source)
        setfield!(o, Symbol("source"), source)
        validate_property(ConnectivityParameters, Symbol("destination"), destination)
        setfield!(o, Symbol("destination"), destination)
        o
    end
end # type ConnectivityParameters

const _property_map_ConnectivityParameters = Dict{Symbol,Symbol}(Symbol("source")=>Symbol("source"), Symbol("destination")=>Symbol("destination"))
const _property_types_ConnectivityParameters = Dict{Symbol,String}(Symbol("source")=>"ConnectivitySource", Symbol("destination")=>"ConnectivityDestination")
Base.propertynames(::Type{ ConnectivityParameters }) = collect(keys(_property_map_ConnectivityParameters))
Swagger.property_type(::Type{ ConnectivityParameters }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_ConnectivityParameters[name]))}
Swagger.field_name(::Type{ ConnectivityParameters }, property_name::Symbol) =  _property_map_ConnectivityParameters[property_name]

function check_required(o::ConnectivityParameters)
    (getproperty(o, Symbol("source")) === nothing) && (return false)
    (getproperty(o, Symbol("destination")) === nothing) && (return false)
    true
end

function validate_property(::Type{ ConnectivityParameters }, name::Symbol, val)
end


text:
# ---
# title: 1413. Minimum Value to Get Positive Step by Step Sum
# id: problem1413
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/minimum-value-to-get-positive-step-by-step-sum/description/>
# hidden: true
# ---
#
# Given an array of integers `nums`, you start with an initial **positive**
# value _startValue_ _._
#
# In each iteration, you calculate the step by step sum of _startValue_  plus
# elements in `nums` (from left to right).
#
# Return the minimum **positive** value of  _startValue_ such that the step by
# step sum is never less than 1.
#
#
#
# **Example 1:**
#
#
#
#     Input: nums = [-3,2,-3,4,2]
#     Output: 5
#     Explanation: If you choose startValue = 4, in the third iteration your step by step sum is less than 1.
#     **step by step sum                 startValue = 4 | startValue = 5 | nums
#     **                  (4 **-3** ) = 1  | (5 **-3** ) = 2    |  -3
#                       (1 **+2** ) = 3  | (2 **+2** ) = 4    |   2
#                       (3 **-3** ) = 0  | (4 **-3** ) = 1    |  -3
#                       (0 **+4** ) = 4  | (1 **+4** ) = 5    |   4
#                       (4 **+2** ) = 6  | (5 **+2** ) = 7    |   2
#
#
# **Example 2:**
#
#
#
#     Input: nums = [1,2]
#     Output: 1
#     Explanation: Minimum start value should be positive.
#
#
# **Example 3:**
#
#
#
#     Input: nums = [1,-2,-3]
#     Output: 5
#
#
#
#
# **Constraints:**
#
#   * `1 <= nums.length <= 100`
#   * `-100 <= nums[i] <= 100`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


text:
using Cowsay
using Documenter

DocMeta.setdocmeta!(Cowsay, :DocTestSetup, :(using Cowsay); recursive=true)

makedocs(;
    modules=[Cowsay],
    authors="Thomas A. Christensen II <25492070+MillironX@users.noreply.github.com> and contributors",
    repo="https://github.com/MillironX/cowsay.jl/blob/{commit}{path}#{line}",
    sitename="cowsay.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://millironx.com/Cowsay.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "Available Cowfiles" => "cows.md",
        "Make your own Cowfile" => "cowfiles.md",
    ],
)

deploydocs(;
    repo="github.com/MillironX/Cowsay.jl",
    devbranch="master",
)


text:
"""
`AbstractBasisSelector` : object specifying how a finite basis is selected from
the infinite space of symmetric polynomials.
"""
abstract type AbstractBasisSelector end

abstract type DownsetBasisSelector <: AbstractBasisSelector end

"""
No constraints on the basis - this selects that largest possible basis
subject to additional constraints baked into the one-particle basis.
In practise this should be equivalent to a naive max-norm basis selection.
"""
struct MaxBasis <: DownsetBasisSelector
end

isadmissible(b, bsel::MaxBasis, basis::OneParticleBasis) = true

"""
`struct SimpleSparseBasis`:

The most basic form of a sparse basis selection, using the total degree.
Only the maximum correlation order and maximum degree may be specified.
This should only be used for testing.
"""
struct SimpleSparseBasis <: DownsetBasisSelector
   maxorder::Int
   maxdeg::Float64
end

# for a one-particle basis function
degree(b::NamedTuple, Bsel::SimpleSparseBasis, basis::OneParticleBasis) =
      degree(b, basis)

# for an ν-correlation basis function
# in this case `bb` should be a Vector of NamedTuples
degree(bb, Bsel::SimpleSparseBasis, basis::OneParticleBasis) =
      length(bb) == 0 ? 0 : sum( degree(b, basis) for b in bb )

isadmissible(b::NamedTuple, Bsel::SimpleSparseBasis, basis::OneParticleBasis) =
      (degree(b, Bsel, basis) <= Bsel.maxdeg)

isadmissible(bb, Bsel::SimpleSparseBasis, basis::OneParticleBasis) =
      (degree(bb, Bsel, basis) <= Bsel.maxdeg) && length(bb) <= Bsel.maxorder

maxorder(Bsel::SimpleSparseBasis) = Bsel.maxorder

"""
`SparseBasis`: probably the standard basis selector enabling weighted degree
functions and varying degree for different correlation orders.
Need to add documentation; for now look at the code for `degree` to see how
the weight and degree dictionaries affect the definition of degree.
"""
struct SparseBasis <: DownsetBasisSelector
   maxorder::Int
   weight::Dict{Symbol, Float64}
   degree::Dict{Any, Float64}
   p::Float64
end

maxorder(Bsel::SparseBasis) = Bsel.maxorder

isadmissible(b::NamedTuple, Bsel::SparseBasis, basis::OneParticleBasis) =
      (degree(b, Bsel, basis) <= _maxdeg(Bsel, 0))

function isadmissible(bb, Bsel::SparseBasis, basis::OneParticleBasis)
   ord = length(bb)
   return (degree(bb, Bsel, basis) <= _maxdeg(Bsel, ord)) && ord <= Bsel.maxorder
end

_maxdeg(Bsel::SparseBasis, ord::Integer) =
      haskey(Bsel.degree, ord) ? Bsel.degree[ord] : Bsel.degree["default"]

# for a one-particle basis function
degree(b::NamedTuple, Deg::SparseBasis, basis::OneParticleBasis) =
      degree(b, basis, Deg.weight)

# for an ν-correlation basis function
# in this case `bb` should be a Vector of NamedTuples
degree(bb, Bsel::SparseBasis, basis::OneParticleBasis)  =  (
      length(bb) == 0 ? 0.0
                      : norm( degree.(bb, Ref(Bsel), Ref(basis)), Bsel.p ) )

"""
`CategoryBasisSelector`: sam as `SparseBasis` but allows in addition to specify
maximum correlation orders for within-category correlations, i.e.,

   `maxorder_dict`
Simple basis selector enabling weighted degree
functions and varying degree for different correlation orders, and additional constraint
   to within-category
Need to add documentation; for now look at the code for `degree` to see how
the weight and degree dictionaries affect the definition of degree.
"""
struct CategoryBasisSelector <: DownsetBasisSelector
   maxorder::Int
   maxorder_dict::Dict{Any, Int}
   isym::Symbol
   weight::Dict{Symbol, Float64}
   weight_cat::Dict{Any, Float64}
   degree::Dict{Any, Float64}
   p::Float64
end

maxorder(Bsel::CategoryBasisSelector) = Bsel.maxorder

maxorder(Bsel::CategoryBasisSelector, category) = Bsel.maxorder_dict[category]

_maxdeg(Bsel::CategoryBasisSelector, ord::Integer) =
         haskey(Bsel.degree, ord) ? Bsel.degree[ord] : Bsel.degree["default"]

isadmissible(b::NamedTuple, Bsel::CategoryBasisSelector, basis::OneParticleBasis) =
      (degree(b, Bsel, basis) <= _maxdeg(Bsel, 0))

function isadmissible(bb, Bsel::CategoryBasisSelector, basis::OneParticleBasis)
   cond_ord = length(bb) <= maxorder(Bsel)
   cond_ord_cats = [
         sum([ getproperty(b, Bsel.isym) == s for b in bb ]) <= maxorder(Bsel, s)
                                 for s in keys(Bsel.maxorder_dict) ]

   return (  (degree(bb, Bsel, basis) <= _maxdeg(Bsel, length(bb)))
             && cond_ord
             && all(cond_ord_cats)
         )
end

# for a one-particle basis function
degree(b::NamedTuple, Bsel::CategoryBasisSelector, basis::OneParticleBasis) =
      degree(b, basis, Bsel.weight) * Bsel.weight_cat[getproperty(b, Bsel.isym)]

# for an ν-correlation basis function
# in this case `bb` should be a Vector of NamedTuples
function degree(bb, Bsel::CategoryBasisSelector, basis::OneParticleBasis)
   if length(bb) == 0
      return 0.0
   else
      return norm( degree.(bb, Ref(Bsel), Ref(basis)), Bsel.p )
   end
end


text:
"""
    reset_aggregated_jumps!(integrator,uprev = nothing; update_jump_params=true)

Reset the state of jump processes and associated solvers following a change
in parameters or such.

Notes
    - `update_jump_params=true` will recalculate the rates stored within any
      MassActionJump that was built from the parameter vector. If the parameter
      vector is unchanged this can safely be set to false to improve performance.
"""

function reset_aggregated_jumps!(integrator,uprev = nothing; update_jump_params=true, kwargs...)
     reset_aggregated_jumps!(integrator,uprev,integrator.opts.callback,
                             update_jump_params=update_jump_params, kwargs...)
     nothing
end

function reset_aggregated_jumps!(integrator,uprev,callback::Nothing; update_jump_params=true, kwargs...)
     nothing
end

function reset_aggregated_jumps!(integrator,uprev,callback::CallbackSet; update_jump_params=true, kwargs...)
    if !isempty(callback.discrete_callbacks)
        reset_aggregated_jumps!(integrator,uprev,callback.discrete_callbacks...,
                                update_jump_params=update_jump_params, kwargs...)
    end
    nothing
end

function reset_aggregated_jumps!(integrator,uprev,cb::DiscreteCallback,cbs...; update_jump_params=true, kwargs...)
    if typeof(cb.condition) <: AbstractSSAJumpAggregator
        maj = cb.condition.ma_jumps
        update_jump_params && using_params(maj) && update_parameters!(cb.condition.ma_jumps,integrator.p; kwargs...)
        cb.condition(cb,integrator.u,integrator.t,integrator)
    end
    reset_aggregated_jumps!(integrator,uprev,cbs...;
                            update_jump_params=update_jump_params, kwargs...)
    nothing
end

function reset_aggregated_jumps!(integrator,uprev,cb::DiscreteCallback; update_jump_params=true, kwargs...)
    if typeof(cb.condition) <: AbstractSSAJumpAggregator
        maj = cb.condition.ma_jumps
        update_jump_params && using_params(maj) && update_parameters!(cb.condition.ma_jumps,integrator.p; kwargs...)
        cb.condition(cb,integrator.u,integrator.t,integrator)
    end
    nothing
end


text:
# Copyright (c) 2020-2021, Bank of Canada
# All rights reserved.

using Base.Broadcast

# custom broadcasting style for TSeries
struct TSeriesStyle{F <: Frequency} <: Broadcast.BroadcastStyle end
frequencyof(::Type{<:TSeriesStyle{F}}) where F <: Frequency = F

Base.Broadcast.BroadcastStyle(::Type{<:TSeries{F}}) where F <: Frequency = TSeriesStyle{F}()
Base.Broadcast.BroadcastStyle(S1::TSeriesStyle, S2::TSeriesStyle) = mixed_freq_error(S1, S2)
Base.Broadcast.BroadcastStyle(S::TSeriesStyle{F}, ::TSeriesStyle{F}) where F <: Frequency = S
Base.Broadcast.BroadcastStyle(S::TSeriesStyle, AS::Broadcast.AbstractArrayStyle{N}) where N = N <= 1 ? S : AS

Base.similar(bc::Broadcast.Broadcasted{<:TSeriesStyle}, ::Type{ElType}) where ElType =
    !isa(bc.axes, Nothing) && length(bc.axes) == 1 ? similar(Vector, ElType, bc.axes) : throw(DimensionMismatch("Cannot create a TSeries with axes $axes."))

# This so far is sufficient to do things like `t .+ 5` and even `t1 .+ t2` when `t1` and `t2` have identical axes.

# now we want to do things like `t .+ v` where `v` is a non-TSeries vector of the same length
# we also want to do `t1 .+ t2` when t1 and t2 are of the same frequency but may have different ranges.
#   in this case the operation will reduce itself to the common range.
# we also want to do `s .= t1 .+ t2` where all three are of the same frequency but may have
#  different axes ranges. In this case the operation on t1 and t2 is done over their common range and
#  the result is stored in s over the same range. If the original range of s is smaller, it gets
# resized to include the broadcasted range. If the original range of s is already larger, then
# values outside the broadcasted range are left unchanged.

# figure out the axes range of the result of the broadcast operation
@inline ts_combine_axes(A, B...) = ts_broadcast_shape(axes(A), ts_combine_axes(B...))
@inline ts_combine_axes(A, B) = ts_broadcast_shape(axes(A), axes(B))
@inline ts_combine_axes(A) = axes(A)

ts_broadcast_shape(shape::Tuple) = shape
ts_broadcast_shape(::Tuple{}, shape::Tuple) = shape
ts_broadcast_shape(shape::Tuple, ::Tuple{}) = shape
function ts_broadcast_shape(shape1::Tuple, shape2::Tuple)
    if length(shape1) > 1 || length(shape2) > 1
        throw(ArgumentError("broadcasting TSeries with ndims > 1."))
    else
        return (mit_common_axes(shape1[1], shape2[1]),)
    end
end

mit_common_axes(a::AbstractRange{<:MIT}, b::AbstractRange{<:MIT}) = mixed_freq_error(a, b)
mit_common_axes(a::AbstractRange{MIT{F}}, b::AbstractRange{MIT{F}}) where F <: Frequency = intersect(a, b)
mit_common_axes(a::AbstractRange{<:MIT}, b::Any) = length(a) == length(b) && first(b) == 1 ? a : throw(DimensionMismatch("Cannot broadcast with $(a) and $b."))
mit_common_axes(a::Any, b::AbstractRange{<:MIT}) = mit_common_axes(b, a)

# given the broadcasted range in `shape`, check that all arguments are of compatible shapes and convert the non-TSeries to
# TSeries with the appropriate range.  This is necessary because the indexing of the broadcast is done using the MIT ranges, so
# plain vectors must be viewed as TSeries.
ts_check_axes(shape, x) = x  # fall back for Number and other things.
@inline ts_check_axes(shape, t::TSeries) =
    shape == axes(t) ? t :
        # if the axes are not identical, we create a "view" into the broadcasted range
        TSeries(first(shape[1]), view(t.values, Int(first(shape[1]) - firstindex(t)) .+ (1:length(shape[1]))))
# For vectors other than TSeries, we create a "view" as a TSeries with the broadcasted range
@inline ts_check_axes(shape, t::AbstractVector) = TSeries(first(shape[1]), view(t, Base.OneTo(length(shape[1]))))
# For nested broadcasted argument, we process the same way recursively.
ts_check_axes(shape, bc::BC) where BC <: Base.Broadcast.Broadcasted = ts_instantiate(bc, shape)

function ts_instantiate(bc::Base.Broadcast.Broadcasted{S}, shape) where S <: Base.Broadcast.BroadcastStyle
    args = map(bc.args) do arg
        ts_check_axes(shape, arg)
    end
    return Base.Broadcast.Broadcasted{S}(bc.f, args, shape)
end

function Base.Broadcast.instantiate(bc::Base.Broadcast.Broadcasted{S,Nothing}) where S <: TSeriesStyle
    shape = ts_combine_axes(bc.args...)
    ts_instantiate(bc, shape)
end

# the following two specializations are necessary in order to be able to have destination on the left of .= that has a different range than the broadcasted result on the right
function Base.Broadcast.instantiate(bc::Base.Broadcast.Broadcasted{S,A}) where {S <: Base.Broadcast.BroadcastStyle,A <: Tuple{<:AbstractRange{<:MIT}}}
    shape = ts_combine_axes(bc.args...)
    I = mit_common_axes(bc.axes[1], shape[1])
    ts_instantiate(bc, (I,))
end

# function Base.Broadcast.instantiate(bc::Base.Broadcast.Broadcasted{S,A}) where {S <: TSeriesStyle,A <: Tuple{<:AbstractRange{<:MIT}}}
#     shape = my_combine_axes(bc.args...)
#     I = _common_axes(bc.axes[1], shape[1])
#     ts_instantiate(bc, (I,))
# end

function Base.Broadcast.instantiate(bc::Base.Broadcast.Broadcasted{S,A}) where {S <: Base.Broadcast.AbstractArrayStyle{0},A <: Tuple{<:AbstractRange{<:MIT}}}
    bc
end

function Base.axes(bc::Base.Broadcast.Broadcasted{<:TSeriesStyle})
    bc.axes === nothing ? ts_combine_axes(bc.args...) : bc.axes
end

function Base.axes(bc::Base.Broadcast.Broadcasted{<:TSeriesStyle}, d::Integer)
    d == 1 ? axes(bc)[1] : Base.OneTo(1)
end

@inline ts_get_index(x::Number, p::MIT) = x
@inline ts_get_index(x::TSeries, p::MIT) = x[p]
@inline ts_get_index(x::Base.Broadcast.Broadcasted, p::MIT) = x[p]

function Base.Broadcast.getindex(bc::Base.Broadcast.Broadcasted, p::MIT)
    args = (ts_get_index(arg, p) for arg in bc.args)
    return bc.f(args...)
end

# this specialization allows for the result to be stored in a TSeries
function Base.copyto!(dest::TSeries, bc::Base.Broadcast.Broadcasted{Nothing})
    bcrng = bc.axes[1]
    drng = eachindex(dest)
    if frequencyof(drng) != frequencyof(bcrng)
        mixed_freq_error(drng, bcrng)
    end
    bc′ = Base.Broadcast.preprocess(dest, bc)
    @simd for I = intersect(bcrng, drng)
        @inbounds dest[I] = bc′[I]
    end
    return dest
end

function Base.Broadcast.dotview(t::TSeries, rng::UnitRange{<:MIT})
    rng ⊆ eachindex(t) ? Base.maybeview(t, rng) : Base.maybeview(resize!(t, union(eachindex(t), rng)), rng)
end


text:
using TensorOperations,SUNRepresentations;

let

d = SUNRepresentations.CGC(Irrep(2,1,0),Irrep(2,1,0),Irrep(2,1,0));
@tensor temp[-1 -2;-3 -4] := d[1 2 -1 -2]*conj(d[1 2 -3 -4])

end