module tmodel

# truncated power law model

export write_grid, write_model, write_lambda, write_dust, Parameters, Grid

# The double dot is because we are now inside the model module, and we want to import the
# constants module, which is part of the enclosing DiskJockey package.
using ..constants

# Write the wavelength sampling file. Only run on setup
function write_lambda(lams::Array{Float64, 1}, basedir::AbstractString)
    fcam = open(basedir * "camera_wavelength_micron.inp", "w")
    nlam = length(lams)
    @printf(fcam, "%d\n", nlam)
    for lam in lams
        @printf(fcam, "%.9e\n", lam) # [microns]
    end

    close(fcam)
end

const eqmirror = true # mirror the grid about the z=0 midplane ?
# if we decide to mirror, then ncells = 1/2 of the true value

# Define a grid object which stores all of these variables
# This will not change for the duration of the run
immutable Grid
    nr::Int
    ntheta::Int
    nphi::Int
    ncells::Int
    # cell edges
    Rs::Vector{Float64}
    Thetas::Vector{Float64}
    Phis::Vector{Float64}
    # cell centers
    rs::Vector{Float64}
    thetas::Vector{Float64}
    phis::Vector{Float64}
end

function Grid(nr::Int, ntheta::Int, r_in::Real, r_out::Real, eqmirror::Bool)
    # Specify a 2D axisymmetric *separable* grid in spherical coordinates:
    # {r, theta, phi}, where theta is angle from zenith, phi is azimuth

    # Number of cells in each dimension
    nphi = 1 # axisymmetric disk
    ncells = nr * ntheta * nphi
    r_in = convert(Float64, r_in) * AU # Inner extent of disk
    r_out = convert(Float64, r_out) * AU # Outer extent of disk

    #Define the cell *walls*
    Rs = logspace(log10(r_in), log10(r_out), nr+1) # [cm] logarithmically spaced

    if eqmirror
        ped = 0.1
        #Thetas = linspace(0, pi/2., ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
        Thetas = pi/2. - (logspace(log10(ped), log10(pi/2. + ped), ntheta+1) - ped)[end:-1:1]
        #Spaced closer near the z=0
    else
        Thetas = linspace(0, pi, ntheta+1)
        # [rad] Angles are internally defined in radians, not degrees
    end

    Phis = Float64[0.0, 0.0] # [rad] cell walls for inactive coordinate

    #Define the cell centers as the average between walls
    rs = 0.5 * (Rs[1:end-1] + Rs[2:end])
    thetas = 0.5 * (Thetas[1:end-1] + Thetas[2:end])
    phis = Float64[0.0]

    return Grid(nr, ntheta, nphi, ncells, Rs, Thetas, Phis, rs, thetas, phis)

end

#This function only needs to be run once, upon setup.
function write_grid(basedir::AbstractString, grid::Grid)
    #amr_grid.inp
    f = open(basedir * "amr_grid.inp", "w")

    #Write the header
    @printf(f, "%d\n", 1) #iformat
    @printf(f, "%d\n", 0) #regular grid (no AMR or Oct-tree)
    @printf(f, "%d\n", 100) #spherical coordiantes
    @printf(f, "%d\n", 0) #gridinfo (none needed for now)
    #incl_r incl_phi incl_z #use this axis?
    @printf(f, "%d %d %d \n", 1, 1, 0) # 2D axisymmetric
    #n_r    n_phi   n_z #number of cells in this dimension
    @printf(f, "%d %d %d \n", grid.nr, grid.ntheta, grid.nphi)

    for R in grid.Rs
        @printf(f, "%.9e\n", R)
    end

    for Theta in grid.Thetas
        @printf(f, "%.9e\n", Theta)
    end

    for Phi in grid.Phis
        @printf(f, "%.9e\n", Phi)
    end

    close(f)
end

#Define the 7 parameters listed in Rosenfeld et al., then write functions to go from those parameters to dust density, velocity, microturbulence.

#Let's try defining a parameters type, an object of which gets passed around.

type Parameters
    M_star::Float64 # [M_sun] stellar mass
    r_in::Float64 # [AU] hard inner radius of the disk
    r_out::Float64 # [AU] radius at which density is depleted by delta
    T_10::Float64 # [K] temperature at 10 AU
    q::Float64 # temperature gradient exponent
    gamma::Float64 # surface temperature gradient exponent
    M_gas::Float64 # [M_Sun] disk mass of gas
    ksi::Float64 # [cm s^{-1}] microturbulence
    dpc::Float64 # [pc] distance to system
    incl::Float64 # [degrees] inclination 0 deg = face on, 90 = edge on.
    PA::Float64 # [degrees] position angle (East of North)
    vel::Float64 # [km/s] systemic velocity (positive is redshift/receeding)
    mu_RA::Float64 # [arcsec] central offset in RA
    mu_DEC::Float64 # [arcsec] central offset in DEC
end

# Typical values, from Rosenfeld et al. 2012, Table 1 for V4046 Sgr
M_star = 1.75 # [M_sun] stellar mass
r_c =  45. # [AU] characteristic radius
T_10 =  115. # [K] temperature at 10 AU
q = 0.63 # temperature gradient exponent
gamma = 1.0 # surface temperature gradient exponent
M_CO = 0.933 # [M_earth] disk mass of CO
ksi = 0.14e5 # [cm s^{-1}] microturbulence
dpc = 73.
incl = 33.5 # [degrees] inclination
PA = 73.
vel = 0.0
mu_RA = 0.0
mu_DEC = 0.0

# Assume all inputs to these functions are in CGS units and in *cylindrical* coordinates.
# Parametric type T allows passing individual Float64 or Vectors.
# Alternate functions accept pars passed around, where pars is in M_star, AU, etc...
function velocity{T}(r::T, M_star::Float64)
    sqrt(G * M_star ./ r)
end
velocity{T}(r::T, pars::Parameters) = velocity(r, pars.M_star * M_sun)

function temperature{T}(r::T, T_10::Float64, q::Float64)
    T_10 * (r ./ (10. * AU)).^(-q)
end
temperature{T}(r::T, pars::Parameters) = temperature(r, pars.T_10, pars.q)

function Hp{T}(r::T, M_star::Float64, T_10::Float64, q::Float64)
    temp = temperature(r, T_10, q)
    sqrt(kB * temp .* r.^3./(mu_gas * m_H * G * M_star))
end
Hp{T}(r::T,  pars::Parameters) = Hp(r, pars.M_star * M_sun, pars.T_10, pars.q)

# Calculate the gas surface density
function Sigma(r::Float64, pars::Parameters)
    r_c = 1. * AU # [cm]
    r_in = pars.r_in * AU
    r_out = pars.r_out * AU

    if r > r_in && r < r_out
        Sigma_c = pars.M_gas * M_sun * (2 - pars.gamma) / (2 * pi * r_c^2 * ((r_out/r_c)^(2 - pars.gamma) - (r_in/r_c)^(2 - pars.gamma)))
        return Sigma_c * (r/r_c)^(-pars.gamma)
    else
        return 0.0
    end
end

# Delivers a gas density in g/cm^3
function rho_gas(r::Float64, z::Float64, pars::Parameters)
    H = Hp(r, pars)
    S = Sigma(r, pars)

    return S/(sqrt(2. * pi) * H) * exp(-0.5 * (z/H)^2)
end

# Now, replace these functions to simply multiply rho_gas by X_12CO/m_12CO, or X_13CO/m_13CO, etc.

n_12CO(r::Float64, z::Float64, pars::Parameters) = number_densities["12CO"] * rho_gas(r, z, pars)

n_13CO(r::Float64, z::Float64, pars::Parameters) = number_densities["13CO"] * rho_gas(r, z, pars)

n_C18O(r::Float64, z::Float64, pars::Parameters) = number_densities["C18O"] * rho_gas(r, z, pars)

# It is realistic to include freezout of CO onto dust grains.
# This is the amount by which the number density of the CO is reduced (X_freeze) relative to
# the nominal value.
function X_freeze(temp::Float64, pars::Parameters)
    # If it's cooler than the freezout temperature, reduce the number density by the given factor
    if temp <= pars.T_freeze
        return pars.X_freeze
    # Otherwise, just keep it as is.
    else
        return 1.0
    end
end

function rho_dust(r::Float64, z::Float64, pars::Parameters)
    nCO = n_CO(r, z, pars) # number of CO molecules per cm^3

    # Convert from nCO to nH2
    nH2 = nCO / 7.e-5 # number density ratio

    # Convert from nH2 (assuming nH2 ~ nGas ) to mGas
    mGas = mu_gas * amu * nH2 # [g]

    # Convert from mGas to mDust using Gas/Dust ratio of 100
    mDust = mGas * 0.01 # [g]

    return mDust
end

# Ksi is microturbulent broadining width in units of km/s. Output of this function
# is in cm/s for RADMC (RADMC manual, eqn 7.12)
function microturbulence(ksi::Float64)
    return ksi * 1.e5 # convert from km/s to cm/s
end

microturbulence(pars::Parameters) = microturbulence(pars.ksi)

function write_model(pars::Parameters, basedir::AbstractString, grid::Grid, species::AbstractString)

    funcs = Dict([("12CO", n_12CO), ("13CO", n_13CO), ("C18O", n_C18O)])
    n_CO = funcs[species]

    # numberdens_co.inp
    fdens = open(basedir * "numberdens_" * molnames[species] * ".inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)

    # gas_velocity.inp
    fvel = open(basedir * "gas_velocity.inp", "w")
    @printf(fvel, "%d\n", 1) #iformat
    @printf(fvel, "%d\n", grid.ncells)

    # gas_temperature.inp
    ftemp = open(basedir * "gas_temperature.inp", "w")
    @printf(ftemp, "%d\n", 1) #iformat
    @printf(ftemp, "%d\n", grid.ncells)

    # microturbulence.inp
    fmicro = open(basedir * "microturbulence.inp", "w")
    @printf(fmicro, "%d\n", 1) #iformat
    @printf(fmicro, "%d\n", grid.ncells)

    # Now, we will need to write the three other files as a function of grid position.
    # Therefore we will do *one* loop over these indices, calculate the required value,
    # and write it to the appropriate file.

    #Looping over the cell centers
    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                @printf(fdens, "%.9e\n", n_CO(r_cyl, z, pars))
                @printf(fvel, "0 0 %.9e\n", velocity(r_cyl, pars))
                @printf(ftemp, "%.9e\n", temperature(r_cyl, pars))
                @printf(fmicro, "%.9e\n", microturbulence(pars))
            end
        end
    end

    close(fdens)
    close(fvel)
    close(ftemp)
    close(fmicro)

end

function write_dust(pars::Parameters, basedir::AbstractString, grid::Grid)
    fdens = open(basedir * "dust_density.inp", "w")
    @printf(fdens, "%d\n", 1) #iformat
    @printf(fdens, "%d\n", grid.ncells)
    @printf(fdens, "%d\n", 1) # number of dust species

    for phi in grid.phis
        for theta in grid.thetas
            for r in grid.rs
                #Convert from spherical to cylindrical coordinates
                z = r * cos(theta)
                r_cyl = r * sin(theta)

                @printf(fdens, "%.9e\n", rho_dust(r_cyl, z, pars))
            end
        end
    end

    close(fdens)
end

end


using Parameters
using DiffEqBase
using ForwardDiff
using Sundials
using ModelingToolkit
abstract type AbstractReactor end
export AbstractReactor

struct Reactor{D,Q} <: AbstractReactor
    domain::D
    ode::ODEProblem
    recommendedsolver::Q
    forwardsensitivities::Bool
end

function Reactor(domain::T,y0::Array{W,1},tspan::Tuple,interfaces::Z=[];p::X=DiffEqBase.NullParameters(),forwardsensitivities=false,forwarddiff=false,modelingtoolkit=false) where {T<:AbstractDomain,W<:Real,Z<:AbstractArray,X}
    dydt(dy::X,y::T,p::V,t::Q) where {X,T,Q,V} = dydtreactor!(dy,y,t,domain,interfaces,p=p)
    jacy!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobiany!(J,y,p,t,domain,interfaces,nothing)
    jacyforwarddiff!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianyforwarddiff!(J,y,p,t,domain,interfaces,nothing)
    jacp!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianp!(J,y,p,t,domain,interfaces,nothing)
    jacpforwarddiff!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianpforwarddiff!(J,y,p,t,domain,interfaces,nothing)

    if (forwardsensitivities || !forwarddiff) && domain isa Union{ConstantTPDomain,ConstantVDomain,ConstantPDomain,ParametrizedTPDomain,ParametrizedVDomain,ParametrizedPDomain,ConstantTVDomain,ParametrizedTConstantVDomain,ConstantTAPhiDomain}
        if !forwardsensitivities
            odefcn = ODEFunction(dydt;jac=jacy!,paramjac=jacp!)
        else
            odefcn = ODEFunction(dydt;paramjac=jacp!)
        end
    else
        odefcn = ODEFunction(dydt;jac=jacyforwarddiff!,paramjac=jacpforwarddiff!)
    end
    if forwardsensitivities
        ode = ODEForwardSensitivityProblem(odefcn,y0,tspan,p)
        recsolver = Sundials.CVODE_BDF(linear_solver=:GMRES)
    else
        ode = ODEProblem(odefcn,y0,tspan,p)
        recsolver  = Sundials.CVODE_BDF()
    end
    if modelingtoolkit
        sys = modelingtoolkitize(ode)
        jac = eval(ModelingToolkit.generate_jacobian(sys)[2])
        if (forwardsensitivities || !forwarddiff) && domain isa Union{ConstantTPDomain,ConstantVDomain,ConstantPDomain,ParametrizedTPDomain,ParametrizedVDomain,ParametrizedPDomain,ConstantTVDomain,ParametrizedTConstantVDomain,ConstantTAPhiDomain}
            odefcn = ODEFunction(dydt;jac=jac,paramjac=jacp!)
        else
            odefcn = ODEFunction(dydt;jac=jac,paramjac=jacpforwarddiff!)
        end
        if forwardsensitivities
            ode = ODEForwardSensitivityProblem(odefcn,y0,tspan,p)
        else
            ode = ODEProblem(odefcn,y0,tspan,p)
        end
    end
    return Reactor(domain,ode,recsolver,forwardsensitivities)
end
function Reactor(domains::T,y0s::W,tspan::W2,interfaces::Z=Tuple(),ps::X=DiffEqBase.NullParameters();forwardsensitivities=false,modelingtoolkit=false) where {T<:Tuple,W<:Tuple,Z,X,W2}
    #adjust indexing
    y0 = zeros(sum(length(y) for y in y0s))
    Nvars = 0
    for domain in domains
        Nvars += domain.indexes[end]
    end
    n = Nvars
    k = 1
    for (j,domain) in enumerate(domains)
        Nspcs = length(domain.phase.species)
        Ntherm = length(domain.indexes) - 2
        for i = 1:6, j = 1:size(domain.rxnarray)[2]
            if domain.rxnarray[i,j] != 0
                domain.rxnarray[i,j] += k-1
            end
        end
        for i in 1:length(domain.constantspeciesinds)
            domain.constantspeciesinds[i] += k-1
        end
        for (thermovar,ind) in domain.thermovariabledict
            domain.thermovariabledict[thermovar] += k-1
        end
        domain.indexes[1] = k
        k += Nspcs
        domain.indexes[2] = k-1
        y0[domain.indexes[1]:domain.indexes[2]] = y0s[j][1:Nspcs]
        for m = 3:2+Ntherm
            domain.indexes[m] = n
            y0[n] = y0s[j][Nspcs+m-2]
            n -= 1
        end
    end

    p = Array{Float64,1}()
    phases = []
    phaseinds = Array{Tuple,1}()
    for (i,domain) in enumerate(domains)
        if domain.phase in phases
            ind = findfirst(phases.==domain.phase)
            domain.parameterindexes[1] = phaseinds[ind][1]
            domain.parameterindexes[2] = phaseinds[ind][2]
            ds = [domains[ind] for ind in findall(x->x.phase==domain.phase,domains)]
            if any(.!isa.(ds,AbstractConstantKDomain)) && any(isa.(ds,AbstractConstantKDomain)) #handle different p formats for the same phase
                for d in ds
                    if !isa(d,AbstractConstantKDomain)
                        p[domain.parameterindexes[1]:domain.parameterindexes[2]] .= d.p
                        break
                    end
                end
                for d in ds
                    if isa(d,AbstractConstantKDomain)
                        if forwardsensitivities == true #error if running into sensitivity bug
                            error("forward sensitivities is not supported for domain combinations that share rate coefficients, but treat them differently ex: ConstantTPDomain and ConstantVDomain")
                        end
                        d.alternativepformat = true
                    end
                end
            end
        else
            domain.parameterindexes[1] = length(p)+1
            domain.parameterindexes[2] = length(p)+length(ps[i])
            push!(phaseinds,(length(p)+1,length(p)+length(ps[i])))
            push!(phases,domain.phase)
            p = vcat(p,ps[i])
        end
    end

    for (i,inter) in enumerate(interfaces)
        if isa(inter, AbstractReactiveInternalInterface)
            ind1 = findfirst(isequal(inter.domain1),domains)
            ind2 = findfirst(isequal(inter.domain2),domains)
            inter.domaininds[1] = ind1
            inter.domaininds[2] = ind2
            inter.parameterindexes[1] = length(p)+1
            inter.parameterindexes[2] = length(p)+length(ps[i+length(domains)])
            inter.rxnarray .= getinterfacereactioninds(inter.domain1,inter.domain2,inter.reactions)
            p = vcat(p,ps[i+length(domains)])
        end
    end

    dydt(dy::X,y::T,p::V,t::Q) where {X,T,Q,V} = dydtreactor!(dy,y,t,domains,interfaces,p=p)
    jacy!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianyforwarddiff!(J,y,p,t,domains,interfaces,nothing)
    jacp!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianpforwarddiff!(J,y,p,t,domains,interfaces,nothing)

    if forwardsensitivities
        odefcn = ODEFunction(dydt;paramjac=jacp!)
        ode = ODEForwardSensitivityProblem(odefcn,y0,tspan,p)
        recsolver = Sundials.CVODE_BDF(linear_solver=:GMRES)
        if modelingtoolkit
            sys = modelingtoolkitize(ode)
            jac = eval(ModelingToolkit.generate_jacobian(sys)[2])
            odefcn = ODEFunction(dydt;jac=jac,paramjac=jacp!)
            ode = ODEForwardSensitivityProblem(odefcn,y0,tspan,p)
        end
    else
        odefcn = ODEFunction(dydt;jac=jacy!,paramjac=jacp!)
        ode = ODEProblem(odefcn,y0,tspan,p)
        recsolver  = Sundials.CVODE_BDF()
        if modelingtoolkit
            sys = modelingtoolkitize(ode)
            jac = eval(ModelingToolkit.generate_jacobian(sys)[2])
            odefcn = ODEFunction(dydt;jac=jac,paramjac=jacp!)
            ode = ODEProblem(odefcn,y0,tspan,p)
        end
    end
    return Reactor(domains,ode,recsolver,forwardsensitivities),y0,p
end
export Reactor

@inline function getrate(rxn::T,cs::Array{W,1},kfs::Array{Q,1},krevs::Array{Q,1}) where {T<:AbstractReaction,Q,W<:Real}
    Nreact = length(rxn.reactantinds)
    Nprod = length(rxn.productinds)
    R = 0.0
    if Nreact == 1
        @fastmath @inbounds R += kfs[rxn.index]*cs[rxn.reactantinds[1]]
    elseif Nreact == 2
        @fastmath @inbounds R += kfs[rxn.index]*cs[rxn.reactantinds[1]]*cs[rxn.reactantinds[2]]
    elseif Nreact == 3
        @fastmath @inbounds R += kfs[rxn.index]*cs[rxn.reactantinds[1]]*cs[rxn.reactantinds[2]]*cs[rxn.reactantinds[3]]
    end

    if Nprod == 1
        @fastmath @inbounds R -= krevs[rxn.index]*cs[rxn.productinds[1]]
    elseif Nprod == 2
        @fastmath @inbounds R -= krevs[rxn.index]*cs[rxn.productinds[1]]*cs[rxn.productinds[2]]
    elseif Nprod == 3
        @fastmath @inbounds R -= krevs[rxn.index]*cs[rxn.productinds[1]]*cs[rxn.productinds[2]]*cs[rxn.productinds[3]]
    end

    return R
end
export getrate

@inline function getrates(rarray,cs,kfs,krevs)
    rts = zeros(length(kfs))
    for i = 1:length(rts)
        if @inbounds rarray[2,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]
        elseif @inbounds rarray[3,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]
        else
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]*cs[rarray[3,i]]
        end
        if @inbounds rarray[5,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]
        elseif @inbounds rarray[6,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]
        else
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]*cs[rarray[6,i]]
        end
        @fastmath R = fR - rR

        rts[i] = R
    end
    return rts
end
export getrates

@inline function addreactionratecontributions!(dydt::Q,rarray::Array{W2,2},cs::W,kfs::Z,krevs::Y) where {Q,Z,Y,T,W,W2}
    @inbounds @simd for i = 1:size(rarray)[2]
        if @inbounds rarray[2,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]
        elseif @inbounds rarray[3,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]
        else
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]*cs[rarray[3,i]]
        end
        if @inbounds rarray[5,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]
        elseif @inbounds rarray[6,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]
        else
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]*cs[rarray[6,i]]
        end
        @fastmath R = fR - rR
        @inbounds @fastmath dydt[rarray[1,i]] -= R
        if @inbounds rarray[2,i] != 0
            @inbounds @fastmath dydt[rarray[2,i]] -= R
            if @inbounds rarray[3,i] != 0
                @inbounds @fastmath dydt[rarray[3,i]] -= R
            end
        end
        @inbounds @fastmath dydt[rarray[4,i]] += R
        if @inbounds rarray[5,i] != 0
            @inbounds @fastmath dydt[rarray[5,i]] += R
            if @inbounds rarray[6,i] != 0
                @inbounds @fastmath dydt[rarray[6,i]] += R
            end
        end
    end
end

@inline function addreactionratecontributions!(dydt::Q,rarray::Array{W2,2},cs::W,kfs::Z,krevs::Y,V) where {Q,Z,Y,T,W,W2}
    @inbounds @simd for i = 1:size(rarray)[2]
        if @inbounds rarray[2,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]
        elseif @inbounds rarray[3,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]
        else
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]*cs[rarray[3,i]]
        end
        if @inbounds rarray[5,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]
        elseif @inbounds rarray[6,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]
        else
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]*cs[rarray[6,i]]
        end
        @fastmath R = (fR - rR)*V
        @inbounds @fastmath dydt[rarray[1,i]] -= R
        if @inbounds rarray[2,i] != 0
            @inbounds @fastmath dydt[rarray[2,i]] -= R
            if @inbounds rarray[3,i] != 0
                @inbounds @fastmath dydt[rarray[3,i]] -= R
            end
        end
        @inbounds @fastmath dydt[rarray[4,i]] += R
        if @inbounds rarray[5,i] != 0
            @inbounds @fastmath dydt[rarray[5,i]] += R
            if @inbounds rarray[6,i] != 0
                @inbounds @fastmath dydt[rarray[6,i]] += R
            end
        end
    end
end
export addreactionratecontributions!

@inline function addreactionratecontributionsforwardreverse!(dydt::Q,rarray::Array{W2,2},cs::W,kfs::Z,krevs::Y,V) where {Q,Z,Y,T,W,W2}
    frts = zeros(length(kfs))
    rrts = zeros(length(kfs))
    rts = zeros(length(kfs))
    @inbounds for i = 1:size(rarray)[2]
        if @inbounds rarray[2,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]
        elseif @inbounds rarray[3,i] == 0
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]
        else
            @inbounds @fastmath fR = kfs[i]*cs[rarray[1,i]]*cs[rarray[2,i]]*cs[rarray[3,i]]
        end
        if @inbounds rarray[5,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]
        elseif @inbounds rarray[6,i] == 0
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]
        else
            @inbounds @fastmath rR = krevs[i]*cs[rarray[4,i]]*cs[rarray[5,i]]*cs[rarray[6,i]]
        end
        @inbounds @fastmath frts[i] = fR*V
        @inbounds @fastmath rrts[i] = rR*V
        @inbounds @fastmath R = frts[i] - rrts[i]
        @inbounds rts[i] = R
        @inbounds @fastmath dydt[rarray[1,i]] -= R
        if @inbounds rarray[2,i] != 0
            @inbounds @fastmath dydt[rarray[2,i]] -= R
            if @inbounds rarray[3,i] != 0
                @inbounds @fastmath dydt[rarray[3,i]] -= R
            end
        end
        @inbounds @fastmath dydt[rarray[4,i]] += R
        if @inbounds rarray[5,i] != 0
            @inbounds @fastmath dydt[rarray[5,i]] += R
            if @inbounds rarray[6,i] != 0
                @inbounds @fastmath dydt[rarray[6,i]] += R
            end
        end
    end
    return rts,frts,rrts
end
export addreactionratecontributionsforwardreverse!

@inline function dydtreactor!(dydt::RC,y::U,t::Z,domain::Q,interfaces::B;p::RV=DiffEqBase.NullParameters(),sensitivity::Bool=true) where {RC,RV,B,Z,U,Q<:AbstractDomain}
    dydt .= 0.0
    ns,cs,T,P,V,C,N,mu,kfs,krevs,Hs,Us,Gs,diffs,Cvave,cpdivR,phi = calcthermo(domain,y,t,p)
    addreactionratecontributions!(dydt,domain.rxnarray,cs,kfs,krevs)
    dydt .*= V
    calcdomainderivatives!(domain,dydt,interfaces;t=t,T=T,P=P,Us=Us,Hs=Hs,V=V,C=C,ns=ns,N=N,Cvave=Cvave)
    return dydt
end
@inline function dydtreactor!(dydt::RC,y::U,t::Z,domains::Q,interfaces::B;p::RV=DiffEqBase.NullParameters(),sensitivity::Bool=true) where {RC,RV,B,Z,U,Q<:Tuple}
    cstot = similar(y)
    cstot .= 0.0
    dydt .= 0.0
    domain = domains[1]
    ns,cs,T,P,V,C,N,mu,kfs,krevs,Hs,Us,Gs,diffs,Cvave,cpdivR,phi = calcthermo(domain,y,t,p)
    vns = Array{Any,1}(undef,length(domains))
    vns[1] = ns
    vcs = Array{Any,1}(undef,length(domains))
    vcs[1] = cs
    cstot[domain.indexes[1]:domain.indexes[2]] = cs
    vT = Array{Any,1}(undef,length(domains))
    vT[1] = T
    vP = Array{Any,1}(undef,length(domains))
    vP[1] = P
    vV = Array{Any,1}(undef,length(domains))
    vV[1] = V
    vC = Array{Any,1}(undef,length(domains))
    vC[1] = C
    vN = Array{Any,1}(undef,length(domains))
    vN[1] = N
    vmu = Array{Any,1}(undef,length(domains))
    vmu[1] = mu
    vkfs = Array{Any,1}(undef,length(domains))
    vkfs[1] = kfs
    vkrevs = Array{Any,1}(undef,length(domains))
    vkrevs[1] = krevs
    vHs = Array{Any,1}(undef,length(domains))
    vHs[1] = Hs
    vUs = Array{Any,1}(undef,length(domains))
    vUs[1] = Us
    vGs = Array{Any,1}(undef,length(domains))
    vGs[1] = Gs
    vdiffs = Array{Any,1}(undef,length(domains))
    vdiffs[1] = diffs
    vCvave = Array{Any,1}(undef,length(domains))
    vCvave[1] = Cvave
    vcpdivR = Array{Any,1}(undef,length(domains))
    vcpdivR[1] = cpdivR
    vphi = Array{Any,1}(undef,length(domains))
    vphi[1] = phi
    addreactionratecontributions!(dydt,domain.rxnarray,cstot,kfs,krevs)
    @views dydt[domain.indexes[1]:domain.indexes[2]] .*= V
    for (i,domain) in enumerate(@views domains[2:end])
        k = i + 1
        vns[k],vcs[k],vT[k],vP[k],vV[k],vC[k],vN[k],vmu[k],vkfs[k],vkrevs[k],vHs[k],vUs[k],vGs[k],vdiffs[k],vCvave[k],vcpdivR[k],vphi[k] = calcthermo(domain,y,t,p)
        cstot[domain.indexes[1]:domain.indexes[2]] .= vcs[k]
        addreactionratecontributions!(dydt,domain.rxnarray,cstot,vkfs[k],vkrevs[k])
        @views dydt[domain.indexes[1]:domain.indexes[2]] .*= vV[k]
    end
    for (i,inter) in enumerate(interfaces)
        if isa(inter,AbstractReactiveInternalInterface)
            evaluate(inter,dydt,domains,vT[inter.domaininds[1]],vT[inter.domaininds[2]],vphi[inter.domaininds[1]],vphi[inter.domaininds[2]],vGs[inter.domaininds[1]],vGs[inter.domaininds[2]],cstot,p)
        end
    end
    for (i,domain) in enumerate(domains)
        calcdomainderivatives!(domain,dydt,interfaces;t=t,T=vT[i],P=vP[i],Us=vUs[i],Hs=vHs[i],V=vV[i],C=vC[i],ns=vns[i],N=vN[i],Cvave=vCvave[i])
    end
    return dydt
end
export dydtreactor!

function jacobianyforwarddiff!(J::Q,y::U,p::W,t::Z,domain::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,V<:AbstractDomain}
    f(dy::X,y::Array{T,1}) where {T<:Real,X} = dydtreactor!(dy,y,t,domain,interfaces;p=p,sensitivity=false)
    ForwardDiff.jacobian!(J,f,zeros(size(y)),y)
end
export jacobianyforwarddiff!

function jacobianyforwarddiff(y::U,p::W,t::Z,domain::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,U<:AbstractArray,W,Z<:Real,V<:AbstractDomain}
    J = zeros(length(y),length(y))
    jacobianyforwarddiff!(J,y,p,t,domain,interfaces,colorvec)
    return J
end
export jacobianyforwarddiff

function jacobianyforwarddiff!(J::Q,y::U,p::W,t::Z,domains::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,V<:Tuple}
    f(dy::X,y::Array{T,1}) where {T<:Real,X} = dydtreactor!(dy,y,t,domains,interfaces;p=p,sensitivity=false)
    ForwardDiff.jacobian!(J,f,zeros(size(y)),y)
end

function jacobianyforwarddiff(y::U,p::W,t::Z,domains::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,V<:Tuple}
    J = zeros(length(y),length(y))
    jacobianyforwarddiff!(J,y,p,t,domains,interfaces,colorvec)
    return J
end
# function jacobiany!(J::Q,y::U,p::W,t::Z,domain::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3<:AbstractArray,Q2<:AbstractArray,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,V<:AbstractDomain}
#     f(y::Array{T,1}) where {T<:Real} = dydtreactor!(y,t,domain,interfaces;p=p,sensitivity=false)
#     forwarddiff_color_jacobian!(J,f,y,colorvec=colorvec)
# end
# function jacobiany!(J::Q,y::U,p::W,t::Z,domain::Q4,interfaces::Q3,colorvec::Q2=nothing) where {Q3<:AbstractArray,Q2,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,Q4<:Union{}}
#     ns,cs,T,P,V,C,N,mu,kfs,krevs,Hs,Us,Gs,diffs,Cvave = calcthermo(domain,y,t,p)
#     jacobiany!(y,t,domain,kfs,krevs,J;zero=true)
# end
# export jacobiany!

function jacobianpforwarddiff!(J::Q,y::U,p::W,t::Z,domain::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,V<:AbstractDomain}
    function f(dy::X,p::Array{T,1}) where {X,T<:Real}
        dydtreactor!(dy,y,t,domain,interfaces;p=p,sensitivity=false)
    end
    dy = zeros(length(y))
    ForwardDiff.jacobian!(J,f,dy,p)
end

function jacobianpforwarddiff!(J::Q,y::U,p::W,t::Z,domains::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,V<:Tuple}
    function f(dy::X,p::Array{T,1}) where {X,T<:Real}
        dydtreactor!(dy,y,t,domains,interfaces;p=p,sensitivity=false)
    end
    dy = zeros(length(y))
    ForwardDiff.jacobian!(J,f,dy,p)
end

export jacobianpforwarddiff!

function jacobianpforwarddiff(y::U,p::W,t::Z,domain::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,U<:AbstractArray,W,Z<:Real,V<:AbstractDomain}
    J = zeros(length(y),length(domain.phase.species)+length(domain.phase.reactions))
    jacobianpforwarddiff!(J,y,p,t,domain,interfaces,colorvec)
end

function jacobianpforwarddiff(y::U,p::W,t::Z,domains::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,V<:Tuple}
    J = zeros(length(y),length(domains.phase.species)+length(domains.phase.reactions))
    jacobianpforwarddiff!(J,y,p,t,domains,interfaces,colorvec)
end

export jacobianpforwarddiff

# function jacobianp!(J::Q,y::U,p::W,t::Z,domain::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3<:AbstractArray,Q2<:AbstractArray,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,V<:AbstractDomain}
#     f(p::Array{T,1}) where {T<:Real} = dydtreactor!(y,domain.t[1],domain,interfaces;p=p,sensitivity=false)
#     forwarddiff_color_jacobian!(J,f,p,colorvec=colorvec)
# end
# function jacobianp!(J::Q,y::U,p::W,t::Z,domain::Q4,interfaces::Q3,colorvec::Q2=nothing) where {Q3<:AbstractArray,Q2,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,Q4<:Union{ConstantTPDomain,ConstantTVDomain}}
#     ns,cs,T,P,V,C,N,mu,kfs,krevs,Hs,Us,Gs,diffs,Cvave = calcthermo(domain,y,t,p)
#     dydt = zeros(length(y))
#     addreactionratecontributions!(dydt,domain.rxnarray,cs,kfs,krevs)
#     @views wV = dydt[domain.indexes[1]:domain.indexes[2]]
#     jacobianp!(domain;cs=cs,V=V,T=T,Us=Us,Cvave=Cvave,N=N,kfs=kfs,krevs=krevs,wV=wV,ratederiv=J)
# end

function jacobiany(y::U,p::W,t::Z,domain::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,U<:AbstractArray,W,Z<:Real,V<:AbstractDomain}
    J = zeros(length(y),length(y))
    jacobiany!(J,y,p,t,domain,interfaces,colorvec)
    return J
end
export jacobiany

function jacobianp(y::U,p::W,t::Z,domain::V,interfaces::Q3,colorvec::Q2=nothing) where {Q3,Q2,U<:AbstractArray,W,Z<:Real,V<:AbstractDomain}
    J = zeros(length(y),length(domain.phase.species)+length(domain.phase.reactions))
    jacobianp!(J,y,p,t,domain,interfaces,colorvec)
    return J
end
export jacobianp

@inline function _spreadreactantpartials!(jac::S,deriv::Float64,rxnarray::Array{Int64,2},rxnind::Int64,ind::Int64) where {S<:AbstractArray}
    @inbounds jac[rxnarray[4,rxnind],ind] += deriv
    if @inbounds rxnarray[5,rxnind] !== 0
        @inbounds jac[rxnarray[5,rxnind],ind] += deriv
        if @inbounds rxnarray[6,rxnind] !== 0
            @inbounds jac[rxnarray[6,rxnind],ind] += deriv
        end
    end
end
@inline function _spreadproductpartials!(jac::S,deriv::Float64,rxnarray::Array{Int64,2},rxnind::Int64,ind::Int64) where {S<:AbstractArray}
    @inbounds jac[rxnarray[1,rxnind],ind] += deriv
    if @inbounds rxnarray[2,rxnind] !== 0
        @inbounds jac[rxnarray[2,rxnind],ind] += deriv
        if @inbounds rxnarray[3,rxnind] !== 0
            @inbounds jac[rxnarray[3,rxnind],ind] += deriv
        end
    end
end

@inline function _jacobianynswrtns!(jac::S,rxnarray::Array{Int64,2},rxnind::Int64,cs::Array{Float64,1},kf::Float64,krev::Float64) where {S<:AbstractArray}
    k=kf
    if rxnarray[2,rxnind] == 0
        deriv = k
        @inbounds jac[rxnarray[1,rxnind],rxnarray[1,rxnind]] -= deriv
        @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[1,rxnind])
    elseif rxnarray[3,rxnind] == 0
        if rxnarray[1,rxnind] == rxnarray[2,rxnind]
            @inbounds @fastmath deriv = 2.0*k*cs[rxnarray[1,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[1,rxnind]] -= 2.0*deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[1,rxnind])
        else
            @inbounds @fastmath deriv = k*cs[rxnarray[2,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[1,rxnind]] -= deriv
            @inbounds jac[rxnarray[2,rxnind],rxnarray[1,rxnind]] -= deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[1,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[1,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[2,rxnind]] -= deriv
            @inbounds jac[rxnarray[2,rxnind],rxnarray[2,rxnind]] -= deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[2,rxnind])
        end
    else
        if rxnarray[1,rxnind]==rxnarray[2,rxnind] && rxnarray[1,rxnind]==rxnarray[3,rxnind]
            @inbounds @fastmath deriv = 3.0*k*cs[rxnarray[1,rxnind]]*cs[rxnarray[1,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[1,rxnind]] -= 3.0*deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[1,rxnind])
        elseif rxnarray[1,rxnind]==rxnarray[2,rxnind]
            @inbounds @fastmath deriv = 2.0*k*cs[rxnarray[1,rxnind]]*cs[rxnarray[3,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[1,rxnind]] -= 2.0*deriv
            @inbounds jac[rxnarray[3,rxnind],rxnarray[1,rxnind]] -= deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[1,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[1,rxnind]]*cs[rxnarray[1,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[3,rxnind]] -= 2.0*deriv
            @inbounds jac[rxnarray[3,rxnind],rxnarray[3,rxnind]] -= deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[3,rxnind])
        elseif rxnarray[2,rxnind]==rxnarray[3,rxnind]
            @inbounds @fastmath deriv = k*cs[rxnarray[2,rxnind]]*cs[rxnarray[2,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[1,rxnind]] -= deriv
            @inbounds jac[rxnarray[2,rxnind],rxnarray[1,rxnind]] -= 2.0*deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[1,rxnind])
            @inbounds @fastmath deriv = 2.0*k*cs[rxnarray[1,rxnind]]*cs[rxnarray[2,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[2,rxnind]] -= deriv
            @inbounds jac[rxnarray[2,rxnind],rxnarray[2,rxnind]] -= 2.0*deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[2,rxnind])
        elseif rxnarray[1,rxnind]==rxnarray[3,rxnind]
            @inbounds @fastmath deriv = 2.0*k*cs[rxnarray[1,rxnind]]*cs[rxnarray[2,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[1,rxnind]] -= 2.0*deriv
            @inbounds jac[rxnarray[2,rxnind],rxnarray[1,rxnind]] -= deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[1,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[1,rxnind]]*cs[rxnarray[1,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[2,rxnind]] -= 2.0*deriv
            @inbounds jac[rxnarray[2,rxnind],rxnarray[2,rxnind]] -= deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[2,rxnind])
        else
            @inbounds @fastmath deriv = k*cs[rxnarray[2,rxnind]]*cs[rxnarray[3,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[1,rxnind]] -= deriv
            @inbounds jac[rxnarray[2,rxnind],rxnarray[1,rxnind]] -= deriv
            @inbounds jac[rxnarray[3,rxnind],rxnarray[1,rxnind]] -= deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[1,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[1,rxnind]]*cs[rxnarray[3,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[2,rxnind]] -= deriv
            @inbounds jac[rxnarray[2,rxnind],rxnarray[2,rxnind]] -= deriv
            @inbounds jac[rxnarray[3,rxnind],rxnarray[2,rxnind]] -= deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[2,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[1,rxnind]]*cs[rxnarray[2,rxnind]]
            @inbounds jac[rxnarray[1,rxnind],rxnarray[3,rxnind]] -= deriv
            @inbounds jac[rxnarray[2,rxnind],rxnarray[3,rxnind]] -= deriv
            @inbounds jac[rxnarray[3,rxnind],rxnarray[3,rxnind]] -= deriv
            @inbounds _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,rxnarray[3,rxnind])
        end
    end
    k=krev
    if rxnarray[5,rxnind] == 0
        deriv = k
        @inbounds jac[rxnarray[4,rxnind],rxnarray[4,rxnind]] -= deriv
        @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[4,rxnind])
    elseif rxnarray[6,rxnind] == 0
        if rxnarray[4,rxnind] == rxnarray[5,rxnind]
            @inbounds @fastmath deriv = 2.0*k*cs[rxnarray[4,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[4,rxnind]] -= 2.0*deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[4,rxnind])
        else
            @inbounds @fastmath deriv = k*cs[rxnarray[5,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[4,rxnind]] -= deriv
            @inbounds jac[rxnarray[5,rxnind],rxnarray[4,rxnind]] -= deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[4,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[4,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[5,rxnind]] -= deriv
            @inbounds jac[rxnarray[5,rxnind],rxnarray[5,rxnind]] -= deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[5,rxnind])
        end
    else
        if rxnarray[4,rxnind]==rxnarray[5,rxnind] && rxnarray[4,rxnind]==rxnarray[6,rxnind]
            @inbounds @fastmath deriv = 3.0*k*cs[rxnarray[4,rxnind]]*cs[rxnarray[4,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[4,rxnind]] -= 3.0*deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[4,rxnind])
        elseif rxnarray[4,rxnind]==rxnarray[5,rxnind]
            @inbounds @fastmath deriv = 2.0*k*cs[rxnarray[4,rxnind]]*cs[rxnarray[6,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[4,rxnind]] -= 2.0*deriv
            @inbounds jac[rxnarray[6,rxnind],rxnarray[4,rxnind]] -= deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[4,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[4,rxnind]]*cs[rxnarray[4,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[6,rxnind]] -= 2.0*deriv
            @inbounds jac[rxnarray[6,rxnind],rxnarray[6,rxnind]] -= deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[6,rxnind])
        elseif rxnarray[5,rxnind]==rxnarray[6,rxnind]
            @inbounds @fastmath deriv = k*cs[rxnarray[5,rxnind]]*cs[rxnarray[5,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[4,rxnind]] -= deriv
            @inbounds jac[rxnarray[5,rxnind],rxnarray[4,rxnind]] -= 2.0*deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[4,rxnind])
            @inbounds @fastmath deriv = 2.0*k*cs[rxnarray[4,rxnind]]*cs[rxnarray[5,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[5,rxnind]] -= deriv
            @inbounds jac[rxnarray[5,rxnind],rxnarray[5,rxnind]] -= 2.0*deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[5,rxnind])
        elseif rxnarray[4,rxnind]==rxnarray[6,rxnind]
            @inbounds @fastmath deriv = 2.0*k*cs[rxnarray[4,rxnind]]*cs[rxnarray[5,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[4,rxnind]] -= 2.0*deriv
            @inbounds jac[rxnarray[5,rxnind],rxnarray[4,rxnind]] -= deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[4,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[4,rxnind]]*cs[rxnarray[4,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[5,rxnind]] -= 2.0*deriv
            @inbounds jac[rxnarray[5,rxnind],rxnarray[5,rxnind]] -= deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[5,rxnind])
        else
            @inbounds @fastmath deriv = k*cs[rxnarray[5,rxnind]]*cs[rxnarray[6,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[4,rxnind]] -= deriv
            @inbounds jac[rxnarray[5,rxnind],rxnarray[4,rxnind]] -= deriv
            @inbounds jac[rxnarray[6,rxnind],rxnarray[4,rxnind]] -= deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[4,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[4,rxnind]]*cs[rxnarray[6,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[5,rxnind]] -= deriv
            @inbounds jac[rxnarray[5,rxnind],rxnarray[5,rxnind]] -= deriv
            @inbounds jac[rxnarray[6,rxnind],rxnarray[5,rxnind]] -= deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[5,rxnind])
            @inbounds @fastmath deriv = k*cs[rxnarray[4,rxnind]]*cs[rxnarray[5,rxnind]]
            @inbounds jac[rxnarray[4,rxnind],rxnarray[6,rxnind]] -= deriv
            @inbounds jac[rxnarray[5,rxnind],rxnarray[6,rxnind]] -= deriv
            @inbounds jac[rxnarray[6,rxnind],rxnarray[6,rxnind]] -= deriv
            @inbounds _spreadproductpartials!(jac,deriv,rxnarray,rxnind,rxnarray[6,rxnind])
        end
    end
end

@inline function _jacobianynswrtV!(jac::S,Vind::Int64,rxnarray::Array{Int64,2},rxnind::Int64,cs::Array{Float64,1},kf::Float64,krev::Float64) where {S<:AbstractArray}
    k=kf
    if rxnarray[2,rxnind]==0
        nothing
    elseif rxnarray[3,rxnind]== 0
        @inbounds @fastmath deriv = -k*cs[rxnarray[1,rxnind]]*cs[rxnarray[2,rxnind]]
        @inbounds jac[rxnarray[1,rxnind],Vind] -= deriv
        @inbounds jac[rxnarray[2,rxnind],Vind] -= deriv
        _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,Vind)
    else
        @inbounds @fastmath deriv = -2.0*k*cs[rxnarray[1,rxnind]]*cs[rxnarray[2,rxnind]]*cs[rxnarray[3,rxnind]]
        @inbounds jac[rxnarray[1,rxnind],Vind] -= deriv
        @inbounds jac[rxnarray[2,rxnind],Vind] -= deriv
        @inbounds jac[rxnarray[3,rxnind],Vind] -= deriv
        _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,Vind)
    end
    k=krev
    if rxnarray[5,rxnind]==0
        nothing
    elseif rxnarray[6,rxnind] == 0
        @inbounds @fastmath deriv = -k*cs[rxnarray[4,rxnind]]*cs[rxnarray[5,rxnind]]
        @inbounds jac[rxnarray[4,rxnind],Vind] -= deriv
        @inbounds jac[rxnarray[5,rxnind],Vind] -= deriv
        _spreadproductpartials!(jac,deriv,rxnarray,rxnind,Vind)
    else
        @inbounds @fastmath deriv = -2.0*k*cs[rxnarray[4,rxnind]]*cs[rxnarray[5,rxnind]]*cs[rxnarray[6,rxnind]]
        @inbounds jac[rxnarray[4,rxnind],Vind] -= deriv
        @inbounds jac[rxnarray[5,rxnind],Vind] -= deriv
        @inbounds jac[rxnarray[6,rxnind],Vind] -= deriv
        _spreadproductpartials!(jac,deriv,rxnarray,rxnind,Vind)
    end
end

"""
This function calculates the ns partials in jacobiany involving k derivatives. dkdx is either dkdni and dkdV. x is either ni or V.
"""
@inline function _jacobianykderiv!(jac::S,xind::Int64,dkfdx::Float64,dkrevdx::Float64,rxnarray::Array{Int64,2},rxnind::Int64,cs::Array{Float64,1},V::Float64) where {S<:AbstractArray}
    dkdx = dkfdx
    if rxnarray[2,rxnind] == 0
        @inbounds @fastmath deriv = dkdx*cs[rxnarray[1,rxnind]]*V
        @inbounds jac[rxnarray[1,rxnind],xind] -= deriv
        _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,xind)
    elseif rxnarray[3,rxnind] == 0
        @inbounds @fastmath deriv = dkdx*cs[rxnarray[1,rxnind]]*cs[rxnarray[2,rxnind]]*V
        @inbounds jac[rxnarray[1,rxnind],xind] -= deriv
        @inbounds jac[rxnarray[2,rxnind],xind] -= deriv
        _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,xind)
    else
        @inbounds @fastmath deriv = dkdx*cs[rxnarray[1,rxnind]]*cs[rxnarray[2,rxnind]]*cs[rxnarray[3,rxnind]]*V
        @inbounds jac[rxnarray[1,rxnind],xind] -= deriv
        @inbounds jac[rxnarray[2,rxnind],xind] -= deriv
        @inbounds jac[rxnarray[3,rxnind],xind] -= deriv
        _spreadreactantpartials!(jac,deriv,rxnarray,rxnind,xind)
    end
    dkdx = dkrevdx
    if rxnarray[5,rxnind] == 0
        @inbounds @fastmath deriv = dkdx*cs[rxnarray[4,rxnind]]*V
        @inbounds jac[rxnarray[4,rxnind],xind] -= deriv
        _spreadproductpartials!(jac,deriv,rxnarray,rxnind,xind)
    elseif rxnarray[6,rxnind] == 0
        @inbounds @fastmath deriv = dkdx*cs[rxnarray[4,rxnind]]*cs[rxnarray[5,rxnind]]*V
        @inbounds jac[rxnarray[4,rxnind],xind] -= deriv
        @inbounds jac[rxnarray[5,rxnind],xind] -= deriv
        _spreadproductpartials!(jac,deriv,rxnarray,rxnind,xind)
    else
        @inbounds @fastmath deriv = dkdx*cs[rxnarray[4,rxnind]]*cs[rxnarray[5,rxnind]]*cs[rxnarray[6,rxnind]]*V
        @inbounds jac[rxnarray[4,rxnind],xind] -= deriv
        @inbounds jac[rxnarray[5,rxnind],xind] -= deriv
        @inbounds jac[rxnarray[6,rxnind],xind] -= deriv
        _spreadproductpartials!(jac,deriv,rxnarray,rxnind,xind)
    end
end

function _dydttherm(dy::X,x::T,y::Q,p::W,t::Z,domain::D,interfaces::Q3,ind::T1) where {X,Q3<:AbstractArray,T<:Real,Q<:AbstractArray,Z<:Real,D<:AbstractDomain,T1<:Integer,W}
    v = [ i != ind ? convert(typeof(x),z) : x for (i,z) in enumerate(y)]
    return dydtreactor!(dy,v,t,domain,interfaces;p=p,sensitivity=false)
end

function jacobianytherm!(jac::Q,y::U,p::W,t::Z,domain::D,interfaces::Q3,ind::I,x::F,colorvec::Q2=nothing) where {Q3<:AbstractArray,Q2,Q<:AbstractArray,U<:AbstractArray,W,Z<:Real,D<:AbstractDomain,I<:Int64,F<:Float64}
    f(dy::X,x::Y) where {Y<:Real,X} = _dydttherm(dy,x,y,p,t,domain,interfaces,ind)
    jac[:,ind] = ForwardDiff.derivative(f,zeros(size(y)),x)
end

# function jacobianp!(d::W;cs::Q,V::Y,T::Y2,Us::Z3,Cvave::Z4,N::Z5,kfs::Z,krevs::X,wV::Q2,ratederiv::Q3) where {Q3,W<:Union{ConstantTPDomain,ConstantTVDomain},Z4<:Real,Z5<:Real,Z3<:AbstractArray,Q2<:AbstractArray,Q<:AbstractArray,Y2<:Real,Y<:Real,Z<:AbstractArray,X<:AbstractArray}
#     Nspcs = length(cs)
#     rxns = d.phase.reactions
#     Nrxns = length(rxns)
#     RTinv = 1.0/(R*T)
#     ratederiv .= 0.0
#
#     for (j,rxn) in enumerate(rxns)
#         Nreact = length(rxn.reactantinds)
#         Nprod = length(rxn.productinds)
#
#         if Nreact == 1
#             rind1 = rxn.reactantinds[1]
#             fderiv = cs[rind1]
#         elseif Nreact == 2
#             rind1,rind2 = rxn.reactantinds
#             fderiv = cs[rind1]*cs[rind2]
#         else
#             rind1,rind2,rind3 = rxn.reactantinds
#             fderiv = cs[rind1]*cs[rind2]*cs[rind3]
#         end
#
#         if Nprod == 1
#             pind1 = rxn.productinds[1]
#             rderiv = krevs[j]/kfs[j]*cs[pind1]
#         elseif Nprod == 2
#             pind1,pind2 = rxn.productinds
#             rderiv = krevs[j]/kfs[j]*cs[pind1]*cs[pind2]
#         else
#             pind1,pind2,pind3 = rxn.productinds
#             rderiv = krevs[j]/kfs[j]*cs[pind1]*cs[pind2]*cs[pind3]
#         end
#
#         flux = fderiv-rderiv
#         gderiv = rderiv*kfs[j]*RTinv
#
#         deriv = zeros(Nspcs)
#
#         deriv[rind1] += gderiv
#         if Nreact > 1
#             deriv[rind2] += gderiv
#             if Nreact > 2
#                 deriv[rind3] == gderiv
#             end
#         end
#
#         deriv[pind1] -= gderiv
#         if Nprod > 1
#             deriv[pind2] -= gderiv
#             if Nprod > 2
#                 deriv[pind3] -= gderiv
#             end
#         end
#
#         ratederiv[rind1,j] -= flux
#         ratederiv[rind1,Nrxns+1:Nrxns+Nspcs] .-= deriv
#         if Nreact > 1
#             ratederiv[rind2,j] -= flux
#             ratederiv[rind2,Nrxns+1:Nrxns+Nspcs] .-= deriv
#             if Nreact > 2
#                 ratederiv[rind3,j] -= flux
#                 ratederiv[rind3,Nrxns+1:Nrxns+Nspcs] .-= deriv
#             end
#         end
#
#         ratederiv[pind1,j] += flux
#         ratederiv[pind1,Nrxns+1:Nrxns+Nspcs] .+= deriv
#         if Nprod > 1
#             ratederiv[pind2,j] += flux
#             ratederiv[pind2,Nrxns+1:Nrxns+Nspcs] .+= deriv
#             if Nprod > 2
#                 ratederiv[pind3,j] += flux
#                 ratederiv[pind3,Nrxns+1:Nrxns+Nspcs] .+= deriv
#             end
#         end
#     end
#     return V*ratederiv
# end
#
# function jacobianp!(d::W; cs::Q,V::Y,T::Y2,Us::Z3,Cvave::Y3,N::Y2,kfs::Z,krevs::X,wV::Q2,ratederiv::Q3) where {W<:Union{ConstantVDomain,ParametrizedVDomain},Q3,Z3<:AbstractArray,Q<:AbstractArray,Q2<:AbstractArray,Y3<:Real,Y2<:Real,Y<:Real,Z<:AbstractArray,X<:AbstractArray}
#     Nspcs = length(cs)
#     rxns = d.phase.reactions
#     Nrxns = length(rxns)
#     RTinv = 1.0/(R*T)
#     ratederiv .= 0.0
#
#     for (j,rxn) in enumerate(rxns)
#         Nreact = length(rxn.reactantinds)
#         Nprod = length(rxn.productinds)
#
#         if Nreact == 1
#             rind1 = rxn.reactantinds[1]
#             fderiv = cs[rind1]
#         elseif Nreact == 2
#             rind1,rind2 = rxn.reactantinds
#             fderiv = cs[rind1]*cs[rind2]
#         else
#             rind1,rind2,rind3 = rxn.reactantinds
#             fderiv = cs[rind1]*cs[rind2]*cs[rind3]
#         end
#
#         if Nprod == 1
#             pind1 = rxn.productinds[1]
#             rderiv = krevs[j]/kfs[j]*cs[pind1]
#         elseif Nprod == 2
#             pind1,pind2 = rxn.productinds
#             rderiv = krevs[j]/kfs[j]*cs[pind1]*cs[pind2]
#         else
#             pind1,pind2,pind3 = rxn.productinds
#             rderiv = krevs[j]/kfs[j]*cs[pind1]*cs[pind2]*cs[pind3]
#         end
#
#         flux = fderiv-rderiv
#         gderiv = rderiv*kfs[j]*RTinv
#
#         deriv = zeros(Nspcs)
#
#         deriv[rind1] += gderiv
#         if Nreact > 1
#             deriv[rind2] += gderiv
#             if Nreact > 2
#                 deriv[rind3] == gderiv
#             end
#         end
#
#         deriv[pind1] -= gderiv
#         if Nprod > 1
#             deriv[pind2] -= gderiv
#             if Nprod > 2
#                 deriv[pind3] -= gderiv
#             end
#         end
#
#         ratederiv[rind1,j] -= flux
#         ratederiv[rind1,Nrxns+1:Nrxns+Nspcs] .-= deriv
#         if Nreact > 1
#             ratederiv[rind2,j] -= flux
#             ratederiv[rind2,Nrxns+1:Nrxns+Nspcs] .-= deriv
#             if Nreact > 2
#                 ratederiv[rind3,j] -= flux
#                 ratederiv[rind3,Nrxns+1:Nrxns+Nspcs] .-= deriv
#             end
#         end
#
#         ratederiv[pind1,j] += flux
#         ratederiv[pind1,Nrxns+1:Nrxns+Nspcs] .+= deriv
#         if Nprod > 1
#             ratederiv[pind2,j] += flux
#             ratederiv[pind2,Nrxns+1:Nrxns+Nspcs] .+= deriv
#             if Nprod > 2
#                 ratederiv[pind3,j] += flux
#                 ratederiv[pind3,Nrxns+1:Nrxns+Nspcs] .+= deriv
#             end
#         end
#     end
#     ratederiv *= V
#     #Temperature stuff
#     @views ratederiv[end,:] += (Us'*ratederiv[1:end-1,:])[1,:]
#     ratederiv[end,1:Nspcs] += wV
#     ratederiv[end,:] /= (N*Cvave)
#     return ratederiv
# end

# @inline function spreadpartials!(jac::S,deriv::T,inds::V,ind::Q,N::Q) where {S<:AbstractArray, T<:Real, V<:AbstractArray, Q<:Integer}
#     if N == 1
#         jac[inds[1],ind] += deriv
#     elseif N == 2
#         jac[inds[1],ind] += deriv
#         jac[inds[2],ind] += deriv
#     elseif N == 3
#         jac[inds[1],ind] += deriv
#         jac[inds[2],ind] += deriv
#         jac[inds[3],ind] += deriv
#     end
# end
#
# @inline function spreadpartials!(jac::S,deriv::T,inds::V,ind::Q,N::Q) where {S<:AbstractArray, T<:Real, V<:AbstractArray, Q<:Integer}
#     if N == 1
#         jac[inds[1],ind] += deriv
#     elseif N == 2
#         jac[inds[1],ind] += deriv
#         jac[inds[2],ind] += deriv
#     elseif N == 3
#         jac[inds[1],ind] += deriv
#         jac[inds[2],ind] += deriv
#         jac[inds[3],ind] += deriv
#     end
# end
#
# function jacobiany!(y::Array{T,1},t::T,domain::ConstantTPDomain,kfs::Array{T,1},krevs::Array{T,1},jac::P;zero::Bool=true) where {P<:AbstractArray,T<:Real,J<:Integer}
#     if zero
#         jac .= 0
#     end
#     N = sum(y)
#     V = N*R*domain.T/domain.P
#     cs = y./V
#     C = N/V
#     rxnarray = domain.rxnarray
#     Nrxns = size(rxnarray)[2]
#     Nspcs = length(y)
#     for i in 1:Nrxns
#         kf = kfs[i]
#         krev = krevs[i]
#         if rxnarray[2,i] == 0
#             jac[rxnarray[1,i],rxnarray[1,i]] -= kf
#             if rxnarray[5,i] == 0
#                 jac[rxnarray[4,i],rxnarray[1,i]] += kf
#             elseif rxnarray[6,i] == 0
#                 jac[rxnarray[4,i],rxnarray[1,i]] += kf
#                 jac[rxnarray[5,i],rxnarray[1,i]] += kf
#             else
#                 jac[rxnarray[4,i],rxnarray[1,i]] += kf
#                 jac[rxnarray[5,i],rxnarray[1,i]] += kf
#                 jac[rxnarray[6,i],rxnarray[1,i]] += kf
#             end
#         elseif rxnarray[3,i] == 0
#             corr = -kf*cs[rxnarray[1,i]]*cs[rxnarray[2,i]]/C #correction for the partial of the volume term
#             if rxnarray[1,i] == rxnarray[2,i]
#                 deriv = 2*kf*cs[rxnarray[1,i]]
#                 jac[rxnarray[1,i],rxnarray[1,i]] -= 2.0*deriv
#                 for j in 1:Nspcs
#                     jac[rxnarray[1,i],j] -= 2.0*corr
#                 end
#                 jac[rxnarray[4,i],rxnarray[1,i]] += deriv
#                 for j in 1:Nspcs
#                     jac[rxnarray[4,i],j] += corr
#                 end
#                 if rxnarray[5,i] != 0
#                     jac[rxnarray[5,i],rxnarray[1,i]] += deriv
#                     for j = 1:Nspcs
#                         jac[rxnarray[5,i],j] += corr
#                     end
#                     if rxnarray[6,i] != 0
#                         jac[rxnarray[6,i],rxnarray[1,i]] += deriv
#                         for j = 1:Nspcs
#                             jac[rxnarray[6,i],j] += corr
#                         end
#                     end
#                 end
#             else
#                 #derivative with respect to reactant 1
#                 deriv = kf*cs[rxnarray[2,i]]
#                 jac[rxnarray[1,i],rxnarray[1,i]] -= deriv
#                 jac[rxnarray[2,i],rxnarray[1,i]] -= deriv
#
#                 jac[rxnarray[4,i],rxnarray[1,i]] += deriv
#                 if rxnarray[5,i] != 0
#                     jac[rxnarray[5,i],rxnarray[1,i]] += deriv
#                     if rxnarray[6,i] != 0
#                         jac[rxnarray[6,i],rxnarray[1,i]] += deriv
#                     end
#                 end
#
#                 #derivative with respect to reactant 2
#                 deriv = kf*cs[rxnarray[1,i]]
#                 jac[rxnarray[1,i],rxnarray[2,i]] -= deriv
#                 jac[rxnarray[2,i],rxnarray[2,i]] -= deriv
#                 for j = 1:Nspcs
#                     jac[rxnarray[1,i],j] -= corr
#                     jac[rxnarray[2,i],j] -= corr
#                 end
#                 jac[rxnarray[4,i],rxnarray[2,j]] += deriv
#                 if rxnarray[5,i] != 0
#                     jac[rxnarray[5,i],rxnarray[2,i]] += deriv
#                     for j = 1:Nspcs
#                         jac[rxnarray[5,i],j] += corr
#                     end
#                     if rxnarray[6,i] != 0
#                         jac[rxnarray[6,i],rxnarray[2,i]] += deriv
#                         for j = 1:Nspcs
#                             jac[rxnarray[6,i],j] += corr
#                         end
#                     end
#                 end
#             end
#         else
#             corr = -2.0*kf*cs[rxnarray[1,i]]*cs[rxnarray[2,i]]*cs[rxnarray[3,i]]/C
#             if (rxnarray[1,i] == rxnarray[2,i] && rxnarray[1,i] == rxnarray[3,i])
#                 deriv = 3.0*kf*cs[rxnarray[1,i]]*cs[rxnarray[1,i]]
#                 jac[rxnarray[1,i],rxnarray[1,i]] -= 3.0*deriv
#                 for j = 1:Nspcs
#                     jac[rxnarray[1,i],j] -= 3.0*corr
#                 end
#                 jac[rxnarray[4,i],rxnarray[1,i]] += deriv
#                 for j = 1:Nspcs
#                     jac[rxnarray[4,i],j] += corr
#                 end
#                 if rxnarray[5,i] != 0
#                     jac[rxnarray[5,i],rxnarray[1,i]] += deriv
#                     for j = 1:Nspcs
#                         jac[rxnarray[5,i],j] += corr
#                     end
#                     if rxnarray[6,i] != 0
#                         jac[rxnarray[6,i],rxnarray[1,i]] += deriv
#                         for j = 1:Nspcs
#                             jac[rxnarray[6,i],j] += corr
#                         end
#                     end
#                 end
#             elseif rxnarray[1,i] == rxnarray[2,i]
#                 #derivative with respect to reactant 1
#                 deriv = 2.0*kf*cs[rxnarray[1,i]]*cs[rxnarray[3,i]]
#                 jac[rxnarray[1,i],rxnarray[1,i]] -= 2.0*deriv
#                 jac[rxnarray[3,i],rxnarray[1,i]] -= deriv
#
#                 jac[rxnarray[4,i],rxnarray[1,i]] += deriv
#                 if rxnarray[5,i] != 0
#                     jac[rxnarray[5,i],rxnarray[1,i]] += deriv
#                     if rxnarray[6,i] != 0
#                         jac[rxnarray[6,i],rxnarray[1,i]] += deriv
#                     end
#                 end
#
#                 #derivative with respect to reactant 3
#                 deriv = kf*cs[rxnarray]
#
#             ind1,ind2,ind3 = rxn.reactantinds
#             corr = -2.0*kf*cs[ind1]*cs[ind2]*cs[ind3]/C
#             deriv = kf*cs[ind1]*cs[ind2]
#             jac[ind1,ind3] -= deriv
#             jac[ind2,ind3] -= deriv
#             jac[ind3,ind3] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind3,Nprod)
#             deriv = kf*cs[ind1]*cs[ind3]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             jac[ind3,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind2,Nprod)
#             deriv = kf*cs[ind3]*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             jac[ind3,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind1,Nprod)
#             for i in 1:length(domain.phase.species)
#                 jac[ind1,i] -= corr
#                 jac[ind2,i] -= corr
#                 jac[ind3,i] -= corr
#                 spreadpartials!(jac,corr,rxn.productinds,i,Nprod)
#             end
#         end
#         #reverse direction
#         if Nprod == 1
#             ind1 = rxn.productinds[1]
#             jac[ind1,ind1] -= krev
#             if Nprod == 1
#                 jac[rxn.reactantinds[1],ind1] += krev
#             elseif Nprod == 2
#                 jac[rxn.reactantinds[1],ind1] += krev
#                 jac[rxn.reactantinds[2],ind1] += krev
#             elseif Nprod == 3
#                 jac[rxn.reactantinds[1],ind1] += krev
#                 jac[rxn.reactantinds[2],ind1] += krev
#                 jac[rxn.reactantinds[3],ind1] += krev
#             end
#         elseif Nprod == 2
#             ind1,ind2 = rxn.productinds
#             corr = -krev*cs[ind1]*cs[ind2]/C #correction for the partial of the volume term
#             deriv = krev*cs[ind1]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind2,Nreact)
#             deriv = kf*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind1,Nreact)
#             for i in 1:length(domain.phase.species)
#                 jac[ind1,i] -= corr
#                 jac[ind2,i] -= corr
#                 spreadpartials!(jac,corr,rxn.reactantinds,i,Nreact)
#             end
#         elseif Nprod == 3
#             ind1,ind2,ind3 = rxn.productinds
#             corr = -2.0*krev*cs[ind1]*cs[ind2]*cs[ind3]/C
#             deriv = krev*cs[ind1]*cs[ind2]
#             jac[ind1,ind3] -= deriv
#             jac[ind2,ind3] -= deriv
#             jac[ind3,ind3] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind3,Nreact)
#             deriv = kf*cs[ind1]*cs[ind3]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             jac[ind3,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind2,Nreact)
#             deriv = kf*cs[ind3]*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             jac[ind3,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind1,Nreact)
#             for i in 1:length(domain.phase.species)
#                 jac[ind1,i] -= corr
#                 jac[ind2,i] -= corr
#                 jac[ind3,i] -= corr
#                 spreadpartials!(jac,corr,rxn.reactantinds,i,Nreact)
#             end
#         end
#     end
#     for ind in domain.constantspeciesinds
#         jac[ind,:] .= 0
#     end
#     return jac
# end
#
# function jacobiany!(y::Array{T,1},t::T,domain::ConstantTPDomain,kfs::Array{T,1},krevs::Array{T,1},jac::P;zero::Bool=true) where {P<:AbstractArray,T<:Real,J<:Integer}
#     if zero
#         jac .= 0
#     end
#     N = sum(y)
#     V = N*R*domain.T/domain.P
#     cs = y./V
#     C = N/V
#     for (i,rxn) in enumerate(domain.phase.reactions)
#         Nreact = length(rxn.reactantinds)
#         Nprod = length(rxn.productinds)
#         kf = kfs[i]
#         krev = krevs[i]
#         if Nreact == 1
#             ind1 = rxn.reactantinds[1]
#             jac[ind1,ind1] -= kf
#             if Nprod == 1
#                 jac[rxn.productinds[1],ind1] += kf
#             elseif Nprod == 2
#                 jac[rxn.productinds[1],ind1] += kf
#                 jac[rxn.productinds[2],ind1] += kf
#             elseif Nprod == 3
#                 jac[rxn.productinds[1],ind1] += kf
#                 jac[rxn.productinds[2],ind1] += kf
#                 jac[rxn.productinds[3],ind1] += kf
#             end
#         elseif Nreact == 2
#             ind1,ind2 = rxn.reactantinds
#             corr = -kf*cs[ind1]*cs[ind2]/C #correction for the partial of the volume term
#             deriv = kf*cs[ind1]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind2,Nprod)
#             deriv = kf*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind1,Nprod)
#             for i in 1:length(domain.phase.species)
#                 jac[ind1,i] -= corr
#                 jac[ind2,i] -= corr
#                 spreadpartials!(jac,corr,rxn.productinds,i,Nprod)
#             end
#         elseif Nreact == 3
#             ind1,ind2,ind3 = rxn.reactantinds
#             corr = -2.0*kf*cs[ind1]*cs[ind2]*cs[ind3]/C
#             deriv = kf*cs[ind1]*cs[ind2]
#             jac[ind1,ind3] -= deriv
#             jac[ind2,ind3] -= deriv
#             jac[ind3,ind3] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind3,Nprod)
#             deriv = kf*cs[ind1]*cs[ind3]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             jac[ind3,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind2,Nprod)
#             deriv = kf*cs[ind3]*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             jac[ind3,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind1,Nprod)
#             for i in 1:length(domain.phase.species)
#                 jac[ind1,i] -= corr
#                 jac[ind2,i] -= corr
#                 jac[ind3,i] -= corr
#                 spreadpartials!(jac,corr,rxn.productinds,i,Nprod)
#             end
#         end
#         #reverse direction
#         if Nprod == 1
#             ind1 = rxn.productinds[1]
#             jac[ind1,ind1] -= krev
#             if Nprod == 1
#                 jac[rxn.reactantinds[1],ind1] += krev
#             elseif Nprod == 2
#                 jac[rxn.reactantinds[1],ind1] += krev
#                 jac[rxn.reactantinds[2],ind1] += krev
#             elseif Nprod == 3
#                 jac[rxn.reactantinds[1],ind1] += krev
#                 jac[rxn.reactantinds[2],ind1] += krev
#                 jac[rxn.reactantinds[3],ind1] += krev
#             end
#         elseif Nprod == 2
#             ind1,ind2 = rxn.productinds
#             corr = -krev*cs[ind1]*cs[ind2]/C #correction for the partial of the volume term
#             deriv = krev*cs[ind1]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind2,Nreact)
#             deriv = kf*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind1,Nreact)
#             for i in 1:length(domain.phase.species)
#                 jac[ind1,i] -= corr
#                 jac[ind2,i] -= corr
#                 spreadpartials!(jac,corr,rxn.reactantinds,i,Nreact)
#             end
#         elseif Nprod == 3
#             ind1,ind2,ind3 = rxn.productinds
#             corr = -2.0*krev*cs[ind1]*cs[ind2]*cs[ind3]/C
#             deriv = krev*cs[ind1]*cs[ind2]
#             jac[ind1,ind3] -= deriv
#             jac[ind2,ind3] -= deriv
#             jac[ind3,ind3] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind3,Nreact)
#             deriv = kf*cs[ind1]*cs[ind3]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             jac[ind3,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind2,Nreact)
#             deriv = kf*cs[ind3]*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             jac[ind3,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind1,Nreact)
#             for i in 1:length(domain.phase.species)
#                 jac[ind1,i] -= corr
#                 jac[ind2,i] -= corr
#                 jac[ind3,i] -= corr
#                 spreadpartials!(jac,corr,rxn.reactantinds,i,Nreact)
#             end
#         end
#     end
#     for ind in domain.constantspeciesinds
#         jac[ind,:] .= 0
#     end
#     return jac
# end
# function jacobiany!(y::Array{T,1},t::T,domain::ConstantTVDomain,kfs::Array{T,1},krevs::Array{T,1},jac::P;zero::Bool=true) where {P<:AbstractArray,T<:Real,J<:Integer}
#     if zero
#         jac .= 0
#     end
#     cs = y./domain.V
#     for (i,rxn) in enumerate(domain.phase.reactions)
#         Nreact = length(rxn.reactantinds)
#         Nprod = length(rxn.productinds)
#         kf = kfs[i]
#         krev = krevs[i]
#         if Nreact == 1
#             ind1 = rxn.reactantinds[1]
#             jac[ind1,ind1] -= kf
#             if Nprod == 1
#                 jac[rxn.productinds[1],ind1] += kf
#             elseif Nprod == 2
#                 jac[rxn.productinds[1],ind1] += kf
#                 jac[rxn.productinds[2],ind1] += kf
#             elseif Nprod == 3
#                 jac[rxn.productinds[1],ind1] += kf
#                 jac[rxn.productinds[2],ind1] += kf
#                 jac[rxn.productinds[3],ind1] += kf
#             end
#         elseif Nreact == 2
#             ind1,ind2 = rxn.reactantinds
#             deriv = kf*cs[ind1]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind2,Nprod)
#             deriv = kf*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind1,Nprod)
#         elseif Nreact == 3
#             ind1,ind2,ind3 = rxn.reactantinds
#             deriv = kf*state.cs[ind1]*cs[ind2]
#             jac[ind1,ind3] -= deriv
#             jac[ind2,ind3] -= deriv
#             jac[ind3,ind3] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind3,Nprod)
#             deriv = kf*cs[ind1]*cs[ind3]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             jac[ind3,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind2,Nprod)
#             deriv = kf*cs[ind3]*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             jac[ind3,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.productinds,ind1,Nprod)
#         end
#         #reverse direction
#         if Nprod == 1
#             ind1 = rxn.productinds[1]
#             jac[ind1,ind1] -= krev
#             if Nprod == 1
#                 jac[rxn.reactantinds[1],ind1] += krev
#             elseif Nprod == 2
#                 jac[rxn.reactantinds[1],ind1] += krev
#                 jac[rxn.reactantinds[2],ind1] += krev
#             elseif Nprod == 3
#                 jac[rxn.reactantinds[1],ind1] += krev
#                 jac[rxn.reactantinds[2],ind1] += krev
#                 jac[rxn.reactantinds[3],ind1] += krev
#             end
#         elseif Nprod == 2
#             ind1,ind2 = rxn.productinds
#             deriv = krev*cs[ind1]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind2,Nreact)
#             deriv = kf*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind1,Nreact)
#         elseif Nprod == 3
#             ind1,ind2,ind3 = rxn.productinds
#             deriv = krev*cs[ind1]*cs[ind2]
#             jac[ind1,ind3] -= deriv
#             jac[ind2,ind3] -= deriv
#             jac[ind3,ind3] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind3,Nreact)
#             deriv = kf*cs[ind1]*cs[ind3]
#             jac[ind1,ind2] -= deriv
#             jac[ind2,ind2] -= deriv
#             jac[ind3,ind2] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind2,Nreact)
#             deriv = kf*cs[ind3]*cs[ind2]
#             jac[ind1,ind1] -= deriv
#             jac[ind2,ind1] -= deriv
#             jac[ind3,ind1] -= deriv
#             spreadpartials!(jac,deriv,rxn.reactantinds,ind1,Nreact)
#         end
#     end
#     for ind in domain.constantspeciesinds
#         jac[ind,:] .= 0
#     end
#     return jac
# end


module TurbulenceClosures

export
  IsotropicDiffusivity,
  ConstantIsotropicDiffusivity,
  ConstantAnisotropicDiffusivity,
  ConstantSmagorinsky,
  SmagorinskyLilly,
  BlasiusSmagorinsky,
  AnisotropicMinimumDissipation,
  RozemaAnisotropicMinimumDissipation,
  VerstappenAnisotropicMinimumDissipation,

  TurbulentDiffusivities,
  calculate_diffusivities!,

  ∇_κ_∇c,
  ∇_κ_∇T,
  ∇_κ_∇S,
  ∂ⱼ_2ν_Σ₁ⱼ,
  ∂ⱼ_2ν_Σ₂ⱼ,
  ∂ⱼ_2ν_Σ₃ⱼ,

  cell_diffusion_timescale

using Oceananigans: @hascuda

@hascuda using CUDAdrv, CUDAnative

using
  Oceananigans,
  Oceananigans.Operators,
  GPUifyLoops

import Oceananigans: with_tracers

using Oceananigans: AbstractArchitecture, AbstractGrid, buoyancy_perturbation, buoyancy_frequency_squared,
                    TracerFields, device, launch_config

####
#### Molecular viscosity and thermal diffusivity definitions
#### Approximate viscosities and thermal diffusivities for seawater at 20ᵒC and 35 psu,
#### according to Sharqawy et al., "Thermophysical properties of seawater: A review of
#### existing correlations and data" (2010).
####

const ν₀ = 1.05e-6
const κ₀ = 1.46e-7

####
#### Abstract types
####

"""
    TurbulenceClosure{FT}

Abstract supertype for turbulence closures with model parameters stored as properties of
type `FT`.
"""
abstract type TurbulenceClosure{FT} end

"""
    IsotropicDiffusivity{FT} <: TurbulenceClosure{FT}

Abstract supertype for turbulence closures that are defined by an isotropic viscosity
and isotropic diffusivities with model parameters stored as properties of type `FT`.
"""
abstract type IsotropicDiffusivity{FT} <: TurbulenceClosure{FT} end

"""
    TensorDiffusivity{FT} <: TurbulenceClosure{FT}

Abstract supertype for turbulence closures that are defined by a tensor viscosity and
tensor diffusivities with model parameters stored as properties of type `FT`.
"""
abstract type TensorDiffusivity{FT} <: TurbulenceClosure{FT} end

"""
    AbstractSmagorinsky{FT}

Abstract supertype for large eddy simulation models based off the model described
by Smagorinsky with model parameters stored as properties of type `FT`.
"""
abstract type AbstractSmagorinsky{FT} <: IsotropicDiffusivity{FT} end

"""
    AbstractAnisotropicMinimumDissipation{FT}

Abstract supertype for large eddy simulation models based on the anisotropic minimum
dissipation principle with model parameters stored as properties of type `FT`.
"""
abstract type AbstractAnisotropicMinimumDissipation{FT} <: IsotropicDiffusivity{FT} end

####
#### Include module code
####

# Fallback constructor for diffusivity types without precomputed diffusivities:
TurbulentDiffusivities(arch::AbstractArchitecture, grid::AbstractGrid, args...) = nothing

include("turbulence_closure_utils.jl")
include("closure_operators.jl")
include("velocity_tracer_gradients.jl")

include("constant_isotropic_diffusivity.jl")
include("constant_anisotropic_diffusivity.jl")
include("smagorinsky.jl")
include("verstappen_anisotropic_minimum_dissipation.jl")
include("rozema_anisotropic_minimum_dissipation.jl")

include("turbulence_closure_diagnostics.jl")

####
#### Some value judgements here
####

"""
    AnisotropicMinimumDissipation

An alias for `VerstappenAnisotropicMinimumDissipation`.
"""
const AnisotropicMinimumDissipation = VerstappenAnisotropicMinimumDissipation

"""
    ConstantSmagorinsky

An alias for `SmagorinskyLilly`.
"""
const ConstantSmagorinsky = SmagorinskyLilly

end # module


@testset "Neuron utility tests" begin
    import NaiveNASflux: neuronutility_safe, neuronutility

    ml(l, lfun=LazyMutable; insize=nin(l)[]) = fluxvertex(l, inputvertex("in", insize, layertype(l)), layerfun = lfun)

    function tr(l, data; loss=Flux.mse)
        example = [(data, 1)];
        Flux.train!((x,y) -> loss(l(x), y), params(l), example, Descent(0.1))
    end

    function tr(l, output, inputs...)
        example = [(inputs, output)];
        Flux.train!((x,y) -> Flux.mse(l(x...), y), params(l), example, Descent(0.1))
    end

    @testset "Utils" begin
        actonly(curr, act, grad) = act

        @testset "Ewma" begin
            m = NaiveNASflux.Ewma(0.3, actonly)
            @test m(missing, [1,2,3,4], :ignored) == [1,2,3,4]
            @test m([1,2,3,4], [5,6,7,8], :ignored) ≈ [3.8, 4.8, 5.8, 6.8]
            @test m(Float32[1,2,3,4], Float32[5,6,7,8], :ignored) ≈ Float32[3.8, 4.8, 5.8, 6.8]
        end

        @testset "NeuronUtilityEvery{3}" begin
            m = NeuronUtilityEvery(3, actonly)
            @test m(:old, :new, :ignored) == :new
            @test m(:old, :new, :ignored) == :old
            @test m(:old, :new, :ignored) == :old
            @test m(:old, :new, :ignored) == :new
            @test m(:old, :new, :ignored) == :old
            @test m(:old, :new, :ignored) == :old
            @test m(:old, :new, :ignored) == :new
            @test m(:old, :new, :ignored) == :old
            @test m(:old, :new, :ignored) == :old
        end
    end

    @testset "Neuron utility Dense default" begin
        l = ml(Dense(3,5))
        @test size(neuronutility(l)) == (5,)
        Δnout!(v -> 1, l => 3)
        @test size(neuronutility(l)) == (8,)
        Δnout!(v -> 1, l => -4)
        @test size(neuronutility(l)) == (4,)
        @test neuronutility(l) ≈ neuronutility_safe(l)
    end

    @testset "Neuron utility Dense default no bias" begin
        l = ml(Dense(ones(5, 3), Flux.Zeros()))
        @test size(neuronutility(l)) == (5,)
        @test neuronutility(l) ≈ neuronutility_safe(l)
    end

    @testset "Neuron utility RNN default" begin
        l = ml(RNN(3,5))
        @test size(neuronutility(l)) == (5,)
        @test neuronutility(l) ≈ neuronutility_safe(l)
    end

    @testset "Neuron utility Conv default" begin
        l = ml(Conv((2,3), 4=>5))
        @test size(neuronutility(l)) == (5,)
        @test neuronutility(l) ≈ neuronutility_safe(l)
    end

    @testset "Neuron utility unkown default" begin
        l = ml(MeanPool((2,2)); insize = 3)
        @test ismissing(neuronutility(l))
        @test neuronutility_safe(l) == 1
    end

    @testset "ActivationContribution no grad" begin
        f(x) = 2 .* x .^ 2
        Flux.Zygote.@nograd f
        l = ml(Dense(2,3), ActivationContribution)
        @test neuronutility(l) == zeros(3)
        tr(l, ones(Float32, 2, 1), loss = f ∘ Flux.mse)
        @test neuronutility(l) == zeros(3)
        @test length(params(l)) == length(params(layer(l)))
    end

    @testset "Functor and trainable" begin
        import NaiveNASflux: weights, bias
        l = ml(Dense(2,3), ActivationContribution)

        neuronutility_org = neuronutility(l)

        @test params(l) == params(layer(l))
        l2 = fmap(x -> x isa AbstractArray ? fill(17, size(x)) : x, l)
        @test unique(neuronutility(l2)) == unique(bias(layer(l2))) == unique(weights(layer(l2))) == [17]
        @test neuronutility_org === neuronutility(l) == zeros(nout(l))
    end

    @testset "Neuron utility Dense act contrib" begin
        l = ml(Dense(3,5), ActivationContribution)
        @test neuronutility(l) == zeros(5)
        tr(l, ones(Float32, 3, 4))
        @test size(neuronutility(l)) == (5,)
        @test length(params(l)) == length(params(layer(l)))
    end

    @testset "Neuron utility Dense act contrib every 4" begin
        l = ml(Dense(3,5), l -> ActivationContribution(l, NeuronUtilityEvery(4)))
        @test neuronutility(l) == zeros(5)
        nvprev = copy(neuronutility(l))
        tr(l, ones(Float32, 3, 4))
        @test neuronutility(l) != nvprev
        nvprev = copy(neuronutility(l))

        tr(l, ones(Float32, 3, 4))
        @test nvprev == neuronutility(l)

        tr(l, ones(Float32, 3, 4))
        tr(l, ones(Float32, 3, 4))
        tr(l, ones(Float32, 3, 4))
        @test nvprev != neuronutility(l)
    end

    @testset "Neuron utility RNN act contrib" begin
        l = ml(RNN(3,5), ActivationContribution)
        @test neuronutility(l) == zeros(5)
        tr(l, ones(Float32, 3, 8))
        @test size(neuronutility(l)) == (5,)
        @test length(params(l)) == length(params(layer(l)))
    end

    @testset "Neuron utility Conv act contrib" begin
        l = ml(Conv((3,3), 2=>5, pad=(1,1)), ActivationContribution)
        @test neuronutility(l) == zeros(5)
        tr(l, ones(Float32, 4,4,2,5))
        @test size(neuronutility(l)) == (5,)
        @test length(params(l)) == length(params(layer(l)))
    end

    @testset "Neuron utility MaxPool act contrib" begin
        l = ml(MaxPool((3,3)), ActivationContribution, insize=2)
        @test ismissing(neuronutility(l))
        tr(l, ones(Float32, 4,4,2,5))

        @test size(neuronutility(l)) == (2,)
    end

    @testset "Neuron utility GlobalMeanPool act contrib" begin
        l = ml(GlobalMeanPool(), ActivationContribution, insize=2)
        @test ismissing(neuronutility(l))
        tr(l, ones(Float32, 4,4,2,5))

        @test size(neuronutility(l)) == (2,)
    end

    @testset "Elem add ActivationContribution" begin
        ac(l) = ActivationContribution(l)
        v = ac >> ml(Dense(2,3)) + ml(Dense(4,3))
        tr(v, [1 1 1]', [1 2 3]', [4 5 6]')
        @test size(neuronutility(v)) == (3,)

        g = CompGraph(vcat(inputs.(inputs(v))...), v)
        @test size(g(ones(Float32, 2,2), ones(Float32, 4, 2))) == (nout(v), 2)
    end

    @testset "Concat ActivationContribution" begin
        v = concat(ml(Dense(2,3)), ml(Dense(4,5)), layerfun=ActivationContribution)
        tr(v,ones(nout(v), 1), [1 2 3]', [4 5 6 7 8]')
        @test size(neuronutility(v)) == (nout(v),)

        g = CompGraph(vcat(inputs.(inputs(v))...), v)
        @test size(g(ones(Float32, 2,2), ones(Float32, 4, 2))) == (nout(v), 2)
    end

    @testset "Function ActivationContribution" begin
        # Not really an absorbvertex, but ActivationContribution should work on stuff which is not FLux layers.
        v = invariantvertex(ActivationContribution(x -> 2x), ml(Dense(2,3)))
        tr(v, ones(nout(v), 1))
        @test size(neuronutility(v)) == (nout(v),)

        g = CompGraph(vcat(inputs.(inputs(v))...), v)
        @test size(g(ones(Float32, 2,2))) == (nout(v), 2)
    end

    @testset "Mutate ActivationContribution" begin
        l = ml(Dense(3,5), ActivationContribution ∘ LazyMutable)
        Δnout!(v -> 1:nout(v), l, -1)
        @test size(l(ones(Float32, 3,2))) == (4, 2)
        @test size(neuronutility(l)) == (4,)
    end

    @testset "Mutate ActivationContribution MaxPool" begin
        l1 = ml(Conv((3,3), 2=>5, pad=(1,1)), ActivationContribution ∘ LazyMutable)
        l2 = fluxvertex(MaxPool((3,3), pad=(1,1)), l1, layerfun = ActivationContribution ∘ LazyMutable)
        g = CompGraph(inputs(l1), l2)

        # Mutate before activation contribution for l2 has been initialized
        Δnout!(v -> 1:nout(v), l1 => 1)
        @test size(neuronutility(l1)) == (6,)
        @test ismissing(neuronutility(l2))

        # This will initialize it
        @test size(g(ones(Float32, 4,4,2,3))) == (2,2,6,3)

        Δnout!(v -> 1:nout(v), l1 => -2)
        tr(g, ones(Float32, 2,2,4,3), ones(Float32, 4,4,2,3))
        @test size(neuronutility(l1)) == size(neuronutility(l2)) == (4,)

        Δnin!(v -> 1:nout(v), l2 => -1)
        tr(g, ones(Float32, 2,2,3,3), ones(Float32, 4,4,2,3))
        @test size(neuronutility(l1)) == size(neuronutility(l2)) == (3,)
    end

    @testset "Add input edge to ActivationContribution concat" begin
        v0 = denseinputvertex("in", 3)
        v1 = fluxvertex("v1", Dense(nout(v0), 4), v0; layerfun=ActivationContribution)
        v2 = fluxvertex("v2", Dense(nout(v0), 3), v0; layerfun=ActivationContribution)
        v3 = concat("v3", v1; layerfun=ActivationContribution)
        v4 = fluxvertex("v4", Dense(nout(v3), 2), v3; layerfun=ActivationContribution)
        v5 = concat("v5", v4, v3, v3; layerfun=ActivationContribution)

        g = CompGraph(v0, v5)
        Flux.gradient(() -> sum(g(ones(Float32, nout(v0), 1))))

        # make sure values have materialized so we don't accidentally have a scalar utility
        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1)

        @test create_edge!(v2, v3)

        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1) + nout(v2)
        @test length(NaiveNASlib.defaultutility(v5)) == nout(v5) == 2 * nout(v3) + nout(v4)

        @test size(g(ones(Float32, nout(v0), 1))) == (nout(v5), 1)
    end

    @testset "Add input edge to ActivationContribution concat fail" begin
        import NaiveNASlib: PostAlign, ΔSizeFailNoOp, FailAlignSizeNoOp
        import NaiveNASflux: neuronutility
        v0 = denseinputvertex("in", 3)
        v1 = fluxvertex("v1", Dense(nout(v0), 4), v0; layerfun=ActivationContribution)
        v2 = fluxvertex("v2", Dense(nout(v0), 3), v0; layerfun=ActivationContribution)
        v3 = concat("v3", v1; layerfun=ActivationContribution)
        v4 = fluxvertex("v4", Dense(nout(v3), 2), v3; layerfun=ActivationContribution)
        v5 = concat("v5", v4, v3, v3; layerfun=ActivationContribution)

        g = CompGraph(v0, v5)
        Flux.gradient(() -> sum(g(ones(Float32, nout(v0), 1))))

        # make sure values have materialized so we don't accidentally have a scalar utility
        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1)

        nvbefore_v3 = copy(neuronutility(v3))
        nvbefore_v5 = copy(neuronutility(v5))

        @test create_edge!(v2, v3; strategy=PostAlign(ΔSizeFailNoOp(), FailAlignSizeNoOp())) == false

        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1)
        @test length(NaiveNASlib.defaultutility(v5)) == nout(v5) == 2 * nout(v3) + nout(v4)

        @test nvbefore_v3 == neuronutility(v3)
        @test nvbefore_v5 == neuronutility(v5)

        @test size(g(ones(Float32, nout(v0), 1))) == (nout(v5), 1)
    end

    @testset "Add input edge to ActivationContribution concat maze" begin
        import NaiveNASlib: PostAlign
        v0 = denseinputvertex("in", 3)
        v1 = fluxvertex("v1", Dense(nout(v0), 4), v0; layerfun=ActivationContribution)
        v2 = fluxvertex("v2", Dense(nout(v0), 3), v0; layerfun=ActivationContribution)
        v3 = concat("v3", v1; layerfun=ActivationContribution)
        v4 = concat("v4", v3; layerfun=ActivationContribution)
        v5 = concat("v5", v3; layerfun=ActivationContribution)
        v6 = concat("v6", v3,v4,v5; layerfun=ActivationContribution)

        g = CompGraph(v0, v6)
        Flux.gradient(() -> sum(g(ones(Float32, nout(v0), 1))))

        # make sure values have materialized so we don't accidentally have a scalar utility
        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1)

        @test create_edge!(v2, v3; strategy=PostAlign())

        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1) + nout(v2)
        @test length(NaiveNASlib.defaultutility(v4)) == nout(v4) == nout(v3)
        @test length(NaiveNASlib.defaultutility(v5)) == nout(v5) == nout(v3)
        @test length(NaiveNASlib.defaultutility(v6)) == nout(v6) == nout(v3) + nout(v4) + nout(v5)

        @test size(g(ones(Float32, nout(v0), 1))) == (nout(v6), 1)
    end

    @testset "Remove input edge to ActivationContribution concat" begin
        v0 = denseinputvertex("in", 3)
        v1 = fluxvertex("v1", Dense(nout(v0), 4), v0; layerfun=ActivationContribution)
        v2 = fluxvertex("v2", Dense(nout(v0), 3), v0; layerfun=ActivationContribution)
        v3 = concat("v3", v1, v2; layerfun=ActivationContribution)
        v4 = fluxvertex("v4", Dense(nout(v3), 2), v3; layerfun=ActivationContribution)
        v5 = concat("v5", v4, v3, v3; layerfun=ActivationContribution)

        g = CompGraph(v0, v5)
        Flux.gradient(() -> sum(g(ones(Float32, nout(v0), 1))))

        # make sure values have materialized so we don't accidentally have a scalar utility
        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1) + nout(v2)

        @test remove_edge!(v2, v3)

        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1)
        @test length(NaiveNASlib.defaultutility(v5)) == nout(v5) == 2 * nout(v3) + nout(v4)

        @test size(g(ones(Float32, nout(v0), 1))) == (nout(v5), 1)
    end

    @testset "Remove input edge to ActivationContribution concat fail" begin
        import NaiveNASlib: PostAlign, ΔSizeFailNoOp, FailAlignSizeNoOp
        import NaiveNASflux: neuronutility

        v0 = denseinputvertex("in", 3)
        v1 = fluxvertex("v1", Dense(nout(v0), 4), v0; layerfun=ActivationContribution)
        v2 = fluxvertex("v2", Dense(nout(v0), 3), v0; layerfun=ActivationContribution)
        v3 = concat("v3", v1, v2; layerfun=ActivationContribution)
        v4 = fluxvertex("v4", Dense(nout(v3), 2), v3; layerfun=ActivationContribution)
        v5 = concat("v5", v4, v3, v3; layerfun=ActivationContribution)

        g = CompGraph(v0, v5)
        Flux.gradient(() -> sum(g(ones(Float32, nout(v0), 1))))

        # make sure values have materialized so we don't accidentally have a scalar utility
        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1) + nout(v2)

        nvbefore_v3 = copy(neuronutility(v3))
        nvbefore_v5 = copy(neuronutility(v5))

        @test remove_edge!(v2, v3; strategy= PostAlign(ΔSizeFailNoOp(), FailAlignSizeNoOp())) == false

        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1) + nout(v2)
        @test length(NaiveNASlib.defaultutility(v5)) == nout(v5) == 2 * nout(v3) + nout(v4)

        @test nvbefore_v3 == neuronutility(v3)
        @test nvbefore_v5 == neuronutility(v5)

        @test size(g(ones(Float32, nout(v0), 1))) == (nout(v5), 1)
    end

    @testset "Remove input edge to ActivationContribution concat maze" begin
        import NaiveNASlib: PostAlign
        v0 = denseinputvertex("in", 3)
        v1 = fluxvertex("v1", Dense(nout(v0), 4), v0; layerfun=ActivationContribution)
        v2 = fluxvertex("v2", Dense(nout(v0), 3), v0; layerfun=ActivationContribution)
        v3 = concat("v3", v1, v2; layerfun=ActivationContribution)
        v4 = concat("v4", v3; layerfun=ActivationContribution)
        v5 = concat("v5", v3; layerfun=ActivationContribution)
        v6 = concat("v6", v3,v4,v5; layerfun=ActivationContribution)

        g = CompGraph(v0, v6)
        Flux.gradient(() -> sum(g(ones(Float32, nout(v0), 1))))

        # make sure values have materialized so we don't accidentally have a scalar utility
        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1) + nout(v2)

        @test remove_edge!(v2, v3; strategy=PostAlign())

        @test length(NaiveNASlib.defaultutility(v3)) == nout(v3) == nout(v1)
        @test length(NaiveNASlib.defaultutility(v4)) == nout(v4) == nout(v3)
        @test length(NaiveNASlib.defaultutility(v5)) == nout(v5) == nout(v3)
        @test length(NaiveNASlib.defaultutility(v6)) == nout(v6) == nout(v3) + nout(v4) + nout(v5)

        @test size(g(ones(Float32, nout(v0), 1))) == (nout(v6), 1)
    end
end