text:
using Manifolds
using ManifoldsBase

using LinearAlgebra
using Distributions
using DoubleFloats
using ForwardDiff
using Random
using ReverseDiff
using StaticArrays
using Test

"""
    test_manifold(m::Manifold, pts::AbstractVector;
        test_forward_diff = true,
        retraction_methods = [],
        inverse_retraction_methods = [])

Tests general properties of manifold `m`, given at least three different points
that lie on it (contained in `pts`).

# Arguments
- `exp_log_atol_multiplier = 1`, change absolute tolerance of exp/log tests
- `inverse_retraction_methods = []`: inverse retraction methods that will be tested.
- `point_distributions = []` : point distributions to test
-  tvector_distributions = []` : tangent vector distributions to test
- `rand_tvector_atol_multiplier = 1` : chage absolute tolerance in testing that
  random tangent vectors are tangent vectors
- `retraction_methods = []`: retraction methods that will be tested.
- `test_forward_diff = true`: if true, automatic differentiation using
  ForwardDiff is tested.
- `test_forward_diff = true`: if true, automatic differentiation using
  ReverseDiff is tested.
- `test_musical_isomorphisms = false` : test musical isomorphisms
- `test_project_tangent = false` : test projections on tangent spaces
- `test_representation_size = true` : test repersentation size of points/tvectprs
- `test_tangent_vector_broadcasting = true` : test boradcasting operators on TangentSpace
- `test_vector_transport = false` : test vector transport
- `test_mutating_rand = false` : test the mutating random function for points on manifolds
"""
function test_manifold(M::Manifold, pts::AbstractVector;
    test_forward_diff = true,
    test_reverse_diff = true,
    test_tangent_vector_broadcasting = true,
    test_project_tangent = false,
    test_representation_size = true,
    test_musical_isomorphisms = false,
    test_vector_transport = false,
    test_mutating_rand = false,
    retraction_methods = [],
    inverse_retraction_methods = [],
    point_distributions = [],
    tvector_distributions = [],
    exp_log_atol_multiplier = 1,
    rand_tvector_atol_multiplier = 1)
    # log/exp
    length(pts) ≥ 3 || error("Not enough points (at least three expected)")
    isapprox(M, pts[1], pts[2]) && error("Points 1 and 2 are equal")
    isapprox(M, pts[1], pts[3]) && error("Points 1 and 3 are equal")

    @testset "dimension" begin
        @test isa(manifold_dimension(M), Integer)
        @test manifold_dimension(M) ≥ 0
        @test manifold_dimension(M) == vector_space_dimension(Manifolds.VectorBundleFibers(Manifolds.TangentSpace, M))
        @test manifold_dimension(M) == vector_space_dimension(Manifolds.VectorBundleFibers(Manifolds.CotangentSpace, M))
    end

    test_representation_size && @testset "representation" begin
        function test_repr(repr)
            @test isa(repr, Tuple)
            for rs ∈ repr
                @test rs > 0
            end
        end

        test_repr(Manifolds.representation_size(M))
        for VS ∈ (Manifolds.TangentSpace, Manifolds.CotangentSpace)
            test_repr(Manifolds.representation_size(Manifolds.VectorBundleFibers(VS, M)))
        end
    end

    @testset "injectivity radius" begin
        @test injectivity_radius(M, pts[1]) > 0
        @test injectivity_radius(M, pts[1]) ≥ injectivity_radius(M)
        for rm ∈ retraction_methods
            @test injectivity_radius(M, pts[1], rm) > 0
            @test injectivity_radius(M, pts[1], rm) ≤ injectivity_radius(M, pts[1])
        end
    end

    tv1 = log(M, pts[1], pts[2])

    @testset "is_manifold_point / is_tangent_vector" begin
        for pt ∈ pts
            @test is_manifold_point(M, pt)
            @test check_manifold_point(M, pt) === nothing
            @test check_manifold_point(M, pt) === nothing
        end
        @test is_tangent_vector(M, pts[1], tv1; atol = eps(eltype(pts[1])))
        @test check_tangent_vector(M, pts[1], tv1; atol = eps(eltype(pts[1]))) === nothing
        @test check_tangent_vector(M, pts[1], tv1; atol = eps(eltype(pts[1]))) === nothing
    end

    @testset "log/exp tests" begin
        tv2 = log(M, pts[2], pts[1])
        @test isapprox(M, pts[2], exp(M, pts[1], tv1))
        @test isapprox(M, pts[1], exp(M, pts[1], tv1, 0))
        @test isapprox(M, pts[2], exp(M, pts[1], tv1, 1))
        @test isapprox(M, pts[1], exp(M, pts[2], tv2))
        @test is_manifold_point(M, retract(M, pts[1], tv1))
        @test isapprox(M, pts[1], retract(M, pts[1], tv1, 0))
        for retr_method ∈ retraction_methods
            @test is_manifold_point(M, retract(M, pts[1], tv1, retr_method))
            @test isapprox(M, pts[1], retract(M, pts[1], tv1, 0, retr_method))
        end
        new_pt = exp(M, pts[1], tv1)
        retract!(M, new_pt, pts[1], tv1)
        @test is_manifold_point(M, new_pt)
        for x ∈ pts
            @test isapprox(M, zero_tangent_vector(M, x), log(M, x, x); atol = eps(eltype(x)) * exp_log_atol_multiplier)
            @test isapprox(M, zero_tangent_vector(M, x), inverse_retract(M, x, x); atol = eps(eltype(x)) * exp_log_atol_multiplier)
            for inv_retr_method ∈ inverse_retraction_methods
                @test isapprox(M, zero_tangent_vector(M, x), inverse_retract(M, x, x, inv_retr_method); atol = eps(eltype(x)) * exp_log_atol_multiplier)
            end
        end
        zero_tangent_vector!(M, tv1, pts[1])
        @test isapprox(M, pts[1], tv1, zero_tangent_vector(M, pts[1]))
        log!(M, tv1, pts[1], pts[2])
        @test norm(M, pts[1], tv1) ≈ sqrt(inner(M, pts[1], tv1, tv1))

        @test isapprox(M, exp(M, pts[1], tv1, 1), pts[2])
        @test isapprox(M, exp(M, pts[1], tv1, 0), pts[1])

        @test distance(M, pts[1], pts[2]) ≈ norm(M, pts[1], tv1)
    end

    @testset "vector spaces tests" begin
        tv = zero_tangent_vector(M, pts[1])
        mts = Manifolds.VectorBundleFibers(Manifolds.TangentSpace, M)
        @test isapprox(M, pts[1], tv, zero_vector(mts, pts[1]))
        zero_vector!(mts, tv, pts[1])
        @test isapprox(M, pts[1], tv, zero_tangent_vector(M, pts[1]))
    end

    @testset "basic linear algebra in tangent space" begin
        @test isapprox(M, pts[1], 0*tv1, zero_tangent_vector(M, pts[1]); atol = eps(eltype(pts[1])))
        @test isapprox(M, pts[1], 2*tv1, tv1+tv1)
        @test isapprox(M, pts[1], 0*tv1, tv1-tv1)
        @test isapprox(M, pts[1], (-1)*tv1, -tv1)
    end

    test_tangent_vector_broadcasting && @testset "broadcasted linear algebra in tangent space" begin
        @test isapprox(M, pts[1], 3*tv1, 2 .* tv1 .+ tv1)
        @test isapprox(M, pts[1], -tv1, tv1 .- 2 .* tv1)
        @test isapprox(M, pts[1], -tv1, .-tv1)
        v = similar(tv1)
        v .= 2 .* tv1 .+ tv1
        @test v ≈ 3*tv1
    end

    test_project_tangent && @testset "project_tangent test" begin
        @test isapprox(M, pts[1], tv1, project_tangent(M, pts[1], tv1))
        tv = similar(tv1)
        project_tangent!(M, tv, pts[1], tv1)
        @test isapprox(M, pts[1], tv, tv1)
    end

    test_vector_transport && @testset "vector transport" begin
        v1 = log(M, pts[1], pts[2])
        v2 = log(M, pts[1], pts[3])
        v1t1 = vector_transport_to(M, pts[1], v1, pts[3])
        v1t2 = vector_transport_direction(M, pts[1], v1, v2)
        @test is_tangent_vector(M, pts[3], v1t1)
        @test is_tangent_vector(M, pts[3], v1t2)
        @test isapprox(M, pts[3], v1t1, v1t2)
        @test isapprox(M, pts[1], vector_transport_to(M, pts[1], v1, pts[1]), v1)
    end

    test_forward_diff && @testset "ForwardDiff support" begin
        exp_f(t) = distance(M, pts[1], exp(M, pts[1], t*tv1))
        d12 = distance(M, pts[1], pts[2])
        for t ∈ 0.1:0.1:0.9
            @test d12 ≈ ForwardDiff.derivative(exp_f, t)
        end

        retract_f(t) = distance(M, pts[1], retract(M, pts[1], t*tv1))
        for t ∈ 0.1:0.1:0.9
            @test ForwardDiff.derivative(retract_f, t) ≥ 0
        end
    end

    test_reverse_diff && @testset "ReverseDiff support" begin
        exp_f(t) = distance(M, pts[1], exp(M, pts[1], t[1]*tv1))
        d12 = distance(M, pts[1], pts[2])
        for t ∈ 0.1:0.1:0.9
            @test d12 ≈ ReverseDiff.gradient(exp_f, [t])[1]
        end

        retract_f(t) = distance(M, pts[1], retract(M, pts[1], t[1]*tv1))
        for t ∈ 0.1:0.1:0.9
            @test ReverseDiff.gradient(retract_f, [t])[1] ≥ 0
        end
    end

    test_musical_isomorphisms && @testset "Musical isomorphisms" begin
        tv_m = log(M, pts[1], pts[2])
        ctv_m = flat(M, pts[1], FVector(TangentSpace, tv_m))
        @test ctv_m.type == CotangentSpace
        tv_m_back = sharp(M, pts[1], ctv_m)
        @test tv_m_back.type == TangentSpace
    end

    @testset "eltype" begin
        tv1 = log(M, pts[1], pts[2])
        @test eltype(tv1) == eltype(pts[1])
        @test eltype(exp(M, pts[1], tv1)) == eltype(pts[1])
    end

    @testset "copyto!" begin
        p2 = similar(pts[1])
        copyto!(p2, pts[2])
        @test isapprox(M, p2, pts[2])

        tv2 = similar(tv1)
        copyto!(tv2, log(M, pts[2], pts[3]))
        @test isapprox(M, pts[2], tv2, log(M, pts[2], pts[3]))
    end

    @testset "point distributions" begin
        prand = similar(pts[1])
        for pd ∈ point_distributions
            for _ in 1:10
                @test is_manifold_point(M, rand(pd))
                if test_mutating_rand
                    rand!(pd, prand)
                    @test is_manifold_point(M, prand)
                end
            end
        end
    end

    @testset "tangent vector distributions" begin
        for tvd ∈ tvector_distributions
            supp = Manifolds.support(tvd)
            for _ in 1:10
                randtv = rand(tvd)
                @test is_tangent_vector(M, supp.x, randtv; atol = rand_tvector_atol_multiplier * eps(eltype(randtv)))
            end
        end
    end
end


text:
"""
$(SIGNATURES)

Write, to the current working directory, a .tex file with the parmater list for the NIPA dataset and parameter values for the TableID parameter.

Arguments
---------
* `b` -- a [`Bea`](@ref) connection

"""
function nipa_metadata_tex(b::Bea)
    url = b.url
    key = b.key
    bea_dataset = b.dataset

    bea_method2 = "GetParameterList"
    query2 = Dict("UserID" => key,
                     "Method" => bea_method2,
                     "DatasetName" => bea_dataset)
    response2 = get(url; query = query2)
    response_json2 = Requests.json(response2)
    paramlist = response_json2["BEAAPI"]["Results"]["Parameter"]

    bea_method3 = "GetParameterValues"
    parameter_name = "TableID"
    query3 = Dict("UserID" => key,
                     "Method" => bea_method3,
                     "DatasetName" => bea_dataset,
                     "ParameterName" => parameter_name)
    response3 = get(url; query = query3)
    response_json3 = Requests.json(response3)
    paramvals = response_json3["BEAAPI"]["Results"]["ParamValue"]

    f = open("NipaMetadata.tex", "w") do f

        # preamble
        write(f, "\\documentclass[12pt]{article} \n")
        write(f, "\\usepackage[margin=1in]{geometry} \n")
        write(f, "\\usepackage{booktabs} \n")
        write(f, "\\usepackage{parskip} \n")
        write(f, "\\usepackage{longtable} \n")
        write(f, "\\usepackage{hyperref} \n")
        write(f, "\\begin{document} \n \n")
        write(f, "Date Created: \\today \n\n")

        # Table of NIPA Dataset parameters
        write(f, "\\section{Dataset Parameters: $bea_dataset}  \n")
        write(f, "\\begin{tabular}{lp{3.5in}} \n")
        write(f, "\\toprule \n")
        for param in paramlist
            write(f, string("Parameter Name: & \\verb|", param["ParameterName"], "| \\\\ \n"))
            for (key, value) in param
                key != "ParameterName" && write(f, string("\\hspace{1em} \\verb|", key, "| & ", value, " \\\\ \n"))
            end
            param != paramlist[end] && write(f, "\\midrule \n")
        end
        write(f, "\\bottomrule \n")
        write(f, "\\end{tabular} \n\n")
        write(f, "\\clearpage \n\n")

        # Table with parameter values for NIPA TableID
        write(f, "\\section{Parameter Values: $bea_dataset $parameter_name}  \n")
        write(f, "\\begin{longtable}{cp{5in}} \n")
        write(f, "\\toprule \n")
        write(f, "\\verb|TableID| & Description \\\\ \n")
        write(f, "\\midrule \\endhead \n")
        for val in paramvals
            write(f, "\\verb|", string(val["TableID"], "| & ", val["Description"], " \\\\ \n"))
        end
        write(f, "\\bottomrule \n")
        write(f, "\\end{longtable} \n\n")

        # end document
        write(f, "\\end{document} \n")
    end

end # function

ext: jl


text:
function realm_align(file)
    filename = file[1][1:end-4]

    img = load(@sprintf("%s.tif", filename))
    img = convert(Array{N0f16}, img)
    img_l = mean(img, dims = 2)

    function img_dist(img1, img2, shift)
        x1 = size(img1)
        x2 = size(img2)
        dist = 0
        for i = (shift+1):x1[1]
            dist += (img1[i] - img2[i-shift])^2
        end
        return sqrt(dist)
    end

    img_1 = img_l[:, 1, 3]
    img_re = zeros(x, y, z)
    @showprogress "Aligning each frame....." for s = 1:z
        min_dist = 1000000
        shift = 0
        if s != 3
            for i = 0:6
                temp_dist = img_dist(img_1, img_l[:, 1, s], i * 9)
                if min_dist > temp_dist
                    min_dist = temp_dist
                    shift = i
                end
            end
        end
        img_re[shift*9+1:end, :, s] = img[1:end-shift*9, :, s]
    end

    img_re = Gray.(convert.(Normed{UInt16,16}, img_re))
    return img_re
end


text:
using ParserCombinator.Parsers.GML
using ZipFile

#io = ZipFile.Reader("gml/polblogs.zip").files[1]
#println("polblogs")
#io = open("gml/polblogs.gml")
#parse_dict(io; debug=false)
#io = open("gml/polblogs.gml")
#@time parse_dict(io; debug=false)

#println("polblogs id")
#io = open("gml/polblogs.gml")
#parse_id_dict(io; debug=false)
#io = open("gml/polblogs.gml")
#@time parse_id_dict(io; debug=false)

println("polblogs")
s = open(readall, "gml/polblogs.gml")
x = parse_dict(s; debug=false)
@test length(x[:graph][1][:edge]) == 19090
@time parse_raw(s; debug=false)


text:
using PaddedMatrices
include(joinpath(pkgdir(PaddedMatrices), "benchmark/blasbench.jl"))

const LIBBLASFEO = Libdl.dlopen("/home/chriselrod/Documents/libraries/blasfeo/lib/libblasfeo.so");
const DGEMM_BLASFEO = Libdl.dlsym(LIBBLASFEO, :blasfeo_dgemm)
const SGEMM_BLASFEO = Libdl.dlsym(LIBBLASFEO, :blasfeo_sgemm)

@inline nottransposed(A) = A
@inline nottransposed(A::Adjoint) = parent(A)
@inline nottransposed(A::Transpose) = parent(A)

for (fm,T) ∈ [(:DGEMM_BLASFEO,Float64), (:SGEMM_BLASFEO, Float32)]
    @eval begin
        function gemmfeo!(C::AbstractMatrix{$T}, A::AbstractMatrix{$T}, B::AbstractMatrix{$T})
            transA = istransposed(A)
            transB = istransposed(B)
            M, N = size(C); K = size(B, 1) % Int32
            pA = nottransposed(A); pB = nottransposed(B)
            ldA = stride(pA, 2) % Int32
            ldB = stride(pB, 2) % Int32
            ldC = stride(C, 2) % Int32
            α = one($T)
            β = zero($T)
            ccall(
                $fm, Cvoid,
                (Ref{UInt8}, Ref{UInt8}, Ref{Int32}, Ref{Int32}, Ref{Int32}, Ref{$T}, Ref{$T},
                 Ref{Int32}, Ref{$T}, Ref{Int32}, Ref{$T}, Ref{$T}, Ref{Int32}),
                transA, transB, (M % Int32), (N % Int32), K, α, pA, ldA, pB, ldB, β, C, ldC
            )
        end
    end
end

function runbenchfeo(::Type{T}, sizes = 2:300) where {T}
    (StructVector ∘ map)(sizes) do sz
        n, k, m = sz, sz, sz
        C1 = Matrix{T}(undef, n, m)
        C2 = similar(C1);
        C3 = similar(C1);
        C4 = similar(C1);
        A  = randa(T, n, k)
        B  = randa(T, k, m)

        jmlt = benchmark_fun!(PaddedMatrices.jmul!, C1, A, B, sz == first(sizes))
        opbt = benchmark_fun!(gemmopenblas!, C2, A, B, sz == first(sizes), C1)
        mklbt= benchmark_fun!(gemmmkl!, C3, A, B, sz == first(sizes), C1)
        bfeot= benchmark_fun!(gemmfeo!, C4, A, B, sz == first(sizes), C1)
        res =  (matrix_size=sz, OpenBLAS=opbt, MKL=mklbt, PaddedMatrices=jmlt, BLASFEO=bfeot)
        @show res
    end
end

istransposed(::PaddedMatrices.AbstractStrideMatrix{<:Any,<:Any,<:Any,1}) = 'N'
istransposed(::PaddedMatrices.AbstractStrideMatrix{<:Any,<:Any,<:Any,1,1}) = 'N'
istransposed(::PaddedMatrices.AbstractStrideMatrix{<:Any,<:Any,<:Any,<:Any,1}) = 'T'
nottransposed(A::PaddedMatrices.AbstractStrideMatrix{<:Any,<:Any,<:Any,1}) = A
nottransposed(A::PaddedMatrices.AbstractStrideMatrix{<:Any,<:Any,<:Any,1,1}) = A
nottransposed(A::PaddedMatrices.AbstractStrideMatrix{<:Any,<:Any,<:Any,<:Any,1}) = A'
using Random
function runstaticbenchfeo(::Type{T}, sizes, f1 = identity, f2 = identity) where {T}
    (StructVector ∘ map)(sizes) do sz
        Apadded = f1(rand!(FixedSizeMatrix{sz,sz,T}(undef)))
        Bpadded = f2(rand!(FixedSizeMatrix{sz,sz,T}(undef)))
        Cpadded1 = FixedSizeMatrix{sz,sz,T}(undef)
        Cpadded2 = FixedSizeMatrix{sz,sz,T}(undef)
        A = f1(Array(nottransposed(Apadded)));
        B = f2(Array(nottransposed(Bpadded)));
        C = similar(A, size(A,1), size(B,2));
        pmt = benchmark_fun!(mul!, Cpadded1, Apadded, Bpadded, sz == first(sizes))
        fmt = benchmark_fun!(gemmfeo!, Cpadded2, Apadded, Bpadded, sz == first(sizes), Cpadded1)
        jmt = benchmark_fun!(PaddedMatrices.jmul!, C, A, B, sz == first(sizes), Cpadded1)
        res = (matrix_size=sz, SizedArray = pmt, BLASFEO=fmt, PaddedMatrices=jmt)
        @show res
    end
end


text:
using PlotlyJS
using RemoteFiles

# For createing plot elements and arrays
PlotElement = GenericTrace{Dict{Symbol,Any}}
PlotArray(; n = 0) = Vector{PlotElement}(undef, n)

colors =
[
    "#1f77b4",  # muted blue
    "#ff7f0e",  # safety orange
    "#2ca02c",  # cooked asparagus green
    "#d62728",  # brick red
    "#9467bd",  # muted purple
    "#8c564b",  # chestnut brown
    "#e377c2",  # raspberry yogurt pink
    "#7f7f7f",  # middle gray
    "#bcbd22",  # curry yellow-green
    "#17becf"   # blue-teal
]

function load_dataframe()
    file_name = "package_requests_by_date.csv.gz"
    stats_base = "https://julialang-logs.s3.amazonaws.com/public_outputs/current"
    r = RemoteFile(stats_base * "/" * file_name, dir = ".")
    download(r)
    DataFrame(load(File(format"CSV", "./package_requests_by_date.csv.gz")))
end

function lookup(name)
    # Look up Registry.toml files in depots
    tomlfiles = String[]
    for d in DEPOT_PATH
        regs = joinpath(d, "registries")
        if isdir(regs)
            for r in readdir(regs)
                toml = joinpath(regs, r, "Registry.toml")
                if isfile(toml)
                    push!(tomlfiles, toml)
                end
            end
        end
    end

    # Look up uuids in toml files
    uuids = Base.UUID[]
    for f in tomlfiles
        toml = TOML.parsefile(f)
        if haskey(toml, "packages")
            for (k, v) in toml["packages"]
                if v["name"] == name
                    push!(uuids, Base.UUID(k))
                end
            end
        end
    end
    return uuids
end

movingaverage(g, n) = [i < n ? mean(g[begin:i]) : mean(g[i-n+1:i]) for i = 1:length(g)]

layout = Layout(;
    title = "",
    width = 600,
    height = 400,
    xaxis = attr(title = "Date"),
    yaxis = attr(title = "Requests"),
)

c_layout = Layout(;
    title = "",
    width = 500,
    height = 400,
    xaxis = attr(title = "Date"),
    yaxis = attr(title = "Cumulative Requests"),
)

function package_dataframe(package_name)
    result = lookup(package_name)
    if length(result) > 0
        uuid = string(result[1])
        mask = ((df."package_uuid" .== uuid) .& (df."status" .== 200))
        return df[mask, :]
    end
    nothing
end

function linear_regression(x, y)
    xd = Dates.value.(x - x[begin])
    mx = mean(xd) # Mean for days
    my = mean(y) # Mean for requests

    β = sum((xd .- mx) .* (y .- my)) / sum((xd .- mx) .* (xd .- mx))
    α = my - β * mx
    return β .* xd .+ α, β, α
end

function scatter_by_client(
    pdf::DataFrame,
    client_type::String,
    data_x,
    data_y,
    check_vals;
    start_date = nothing,
    end_date = nothing,
    name = "",
    color="green"
)
    mask = (pdf."client_type" .== client_type)
    if start_date !== nothing
        sd = Date(start_date[1:10])
        cond = (pdf."date" .>= sd)
        mask = (mask .& cond)
    end
    if end_date !== nothing
        ed = Date(end_date[1:10])
        cond = (pdf."date" .<= ed)
        mask = (mask .& cond)
    end
    cpdf = pdf[mask, :]
    x = cpdf[:, data_x]
    y = cpdf[:, data_y]
    days = Dates.value(x[end] - x[begin])
    total_requests = sum(y)
    # Moving averages
    my = movingaverage(y, 7)
    ly, β, α = linear_regression(x, y)

    precision = 4
    α = round(α, digits = precision)
    β = round(β, digits = precision)
    μ = round(sum(y) / days, digits = precision)

    info = Dict("total" => total_requests, "interval" => days, "α" => α, "β" => β)

    extra_plots = Dict(
        "MA" => scatter(x = x, y = my, mode = "lines", name = "μ", showlegend = false, line_color=color, line_width=1.5),
        "RL" => scatter(x = x, y = ly, mode = "lines", name = "L", showlegend = false, line_color=color, line_width=2),
    )
    plots = [scatter(x = x, y = y, mode = "lines+markers", name = name, line_color=color,line_width=2)]
    append!(plots, [extra_plots[k] for k in check_vals])
    c_plots = [scatter(x = x, y = cumsum(y), mode = "lines", name = name, line_color=color)]
    return plots, c_plots, info
end

function plot_default(;cumulative = false)
    fig = scatter(x = [0], y = [0], mode = "lines")
    lo = layout
    if cumulative
        lo = c_layout
    end
    plot(fig, lo)
end

function plot_graphs(options, check_vals, client_type; start_date = nothing, end_date = nothing)
    if (options !== nothing)
        if length(options) == 0
            return plot_default()
        end
        plot_data = PlotArray()
        c_plot_data = PlotArray()
        stats_info = []

        for (i,op) in enumerate(options)
            plots, cu, info = scatter_by_client(
                package_dataframe(op),
                client_type,
                "date",
                "request_count",
                check_vals,
                start_date = start_date,
                end_date = end_date,
                name = op,
                color=colors[i]
            )
            append!(plot_data, plots)
            append!(c_plot_data, cu)
            append!(stats_info, [info])
        end
        return plot(plot_data, layout), plot(c_plot_data, c_layout), stats_info
    end
    return plot_default(), plot_default(cumulative=true), []
end


text:
trunc(x) = x < 0 ? zero(x) : x
function sqrt_sin(x)
y = trunc(x)
return sin(sqrt(y)+1)
end

max_stars_repo_path: Module 2/Chapter04/sqrt_sin.jl

max_issues_repo_path: Module 2/Chapter04/sqrt_sin.jl

max_forks_repo_path: Module 2/Chapter04/sqrt_sin.jl


text:
module HarwellRutherfordBoeing

using Compat

export HarwellBoeingMatrix, RutherfordBoeingData

include("harwell_boeing.jl")
include("rutherford_boeing.jl")

end


text:
function Draw(nodenum,Node,P)
    xaex = Vector{Int16}()
    for i = 1:nodenum
        if Node[i,2] == 0 && Node[i,3] == 1
            point = append!(xaex,i)
        end
    end

    sort = Vector{Float64}()
    for i = xaex
        sort = append!(sort,Node[i,1])
    end
    println(xaex)
    index = sortperm(sort)
    println(index)

    x = Vector{Float64}()
    for i = index
        x = append!(x,Node[xaex[i],1])
    end
    y = Vector{Float64}()
    for i = index
        y = append!(y,P[xaex[i],3])
    end

    return x,y
end

max_stars_repo_path: src/single-core juSFEM/Draw.jl

max_stars_repo_name: Elliothuo/juSFEM

max_issues_repo_path: src/single-core juSFEM/Draw.jl

max_issues_repo_name: Elliothuo/juSFEM

max_forks_repo_path: src/single-core juSFEM/Draw.jl

max_forks_repo_name: Elliothuo/juSFEM


text:
## Load the packages.
using LoopVectorization
using Interpolations
using Statistics
using BenchmarkTools
using PyPlot
using Printf

## Define the upsample functions.
"""
function upsample_nn!(hi, lo, itot, jtot, ktot, ifac, jfac, kfac, ioff, joff, koff)
    @tturbo for k in 0:ktot-1, j in 0:jtot-1, i in 0:itot-1
        i_lo = i + 2; j_lo = j + 2; k_lo = k + 2
        i_hi = ifac*i + 2; j_hi = jfac*j + 2; k_hi = kfac*k + 2
        for kk in 0:kfac-1, jj in 0:jfac-1, ii in 0:ifac-1
            iii = floor(Int, ioff + 1/ifac * ii)
            jjj = floor(Int, joff + 1/jfac * jj)
            kkk = floor(Int, koff + 1/kfac * kk)
            hi[i_hi+ii, j_hi+jj, k_hi+kk] = lo[i_lo+iii, j_lo+jjj, k_lo+kkk]
        end
    end
end

function upsample_nn_ref!(hi, lo, x_hi, y_hi, z_hi, x_lo, y_lo, z_lo)
    interp = interpolate((x_lo, y_lo, z_lo), lo, (Gridded(Constant()), Gridded(Constant()), Gridded(Constant())))
    hi[2:end-1, 2:end-1, 2:end-1] .= interp(x_hi[2:end-1], y_hi[2:end-1], z_hi[2:end-1])
end
"""

macro upsample_lin_fast(suffix, ifac, jfac, kfac, ioff, joff, koff)
    ex_inner = []
    for kk in 0:kfac-1, jj in 0:jfac-1, ii in 0:ifac-1
        push!(ex_inner, :(i_hi = $ifac*i+$ii+is_hi), :(j_hi = $jfac*j+$jj+js_hi), :(k_hi = $kfac*k+$kk+ks_hi))

        i_pos = -1/2 + ioff + (1/2 - ioff + ii)/ifac
        j_pos = -1/2 + joff + (1/2 - joff + jj)/jfac
        k_pos = -1/2 + koff + (1/2 - koff + kk)/kfac

        i_lo_off = floor(Int, i_pos) + is_lo
        j_lo_off = floor(Int, j_pos) + js_lo
        k_lo_off = floor(Int, k_pos) + ks_lo
        push!(ex_inner, :(i_lo = i + $i_lo_off), :(j_lo = j + $j_lo_off), :(k_lo = k + $k_lo_off))

        fi = mod(i_pos, 1)
        fj = mod(j_pos, 1)
        fk = mod(k_pos, 1)

        coef = zeros(2, 2, 2)
        coef[1, 1, 1] = (((1-fi)*fj+fi-1)*fk+(fi-1)*fj-fi+1)
        coef[2, 1, 1] = ((fi*fj-fi)*fk-fi*fj+fi)
        coef[1, 2, 1] = ((fi-1)*fj*fk+(1-fi)*fj)
        coef[2, 2, 1] = (fi*fj-fi*fj*fk)
        coef[1, 1, 2] = ((fi-1)*fj-fi+1)*fk
        coef[2, 1, 2] = (fi-fi*fj)*fk
        coef[1, 2, 2] = (1-fi)*fj*fk
        coef[2, 2, 2] = fi*fj*fk

        rhs_list = []
        for kc in 1:2, jc in 1:2, ic in 1:2
            if !(coef[ic, jc, kc] ≈ 0)
                push!(rhs_list, :( $(coef[ic, jc, kc])*lo[i_lo+$(ic-1), j_lo+$(jc-1), k_lo+$(kc-1)] ))
            end
        end

        rhs = Expr(:call, :+, rhs_list...)

        push!(ex_inner, :(hi[i_hi, j_hi, k_hi] = $rhs))
    end

    ex_inner_block = Expr(:block, ex_inner...)

    name = Symbol(@sprintf "upsample_lin_%s!" suffix)
    ex = quote
        function $name(hi, lo, itot, jtot, ktot, is_hi, js_hi, ks_hi, is_lo, js_lo, ks_lo)
            Threads.@threads for k in 0:ktot-1
                for j in 0:jtot-1
                    @inbounds @simd for i in 0:itot-1
                        $ex_inner_block
                    end
                end
            end
        end
    end
    return esc(ex)
end

function upsample_lin!(
        hi, lo,
        itot, jtot, ktot,
        is_hi, js_hi, ks_hi, is_lo, js_lo, ks_lo,
        ifac, jfac, kfac, ioff, joff, koff)

    Threads.@threads for k in 0:ktot-1
        for j in 0:jtot-1
            @simd for i in 0:itot-1
                for kk in 0:kfac-1, jj in 0:jfac-1, ii in 0:ifac-1
                    i_hi = ifac*i+ii+is_hi; j_hi = jfac*j+jj+js_hi; k_hi = kfac*k+kk+ks_hi;

                    # Determine fractional distance and west, south, and bot point.
                    i_pos = -1/2 + ioff + (1/2 - ioff + ii)/ifac
                    j_pos = -1/2 + joff + (1/2 - joff + jj)/jfac
                    k_pos = -1/2 + koff + (1/2 - koff + kk)/kfac

                    fi = mod(i_pos, 1)
                    fj = mod(j_pos, 1)
                    fk = mod(k_pos, 1)

                    coef_111 = (((1-fi)*fj+fi-1)*fk+(fi-1)*fj-fi+1)
                    coef_211 = ((fi*fj-fi)*fk-fi*fj+fi)
                    coef_121 = ((fi-1)*fj*fk+(1-fi)*fj)
                    coef_221 = (fi*fj-fi*fj*fk)
                    coef_112 = ((fi-1)*fj-fi+1)*fk
                    coef_212 = (fi-fi*fj)*fk
                    coef_122 = (1-fi)*fj*fk
                    coef_222 = fi*fj*fk

                    i_lo = i + floor(Int, i_pos) + is_lo
                    j_lo = j + floor(Int, j_pos) + js_lo
                    k_lo = k + floor(Int, k_pos) + ks_lo

                    @inbounds hi[i_hi, j_hi, k_hi] = (
                        + coef_111 * lo[i_lo  , j_lo  , k_lo  ]
                        + coef_211 * lo[i_lo+1, j_lo  , k_lo  ]
                        + coef_121 * lo[i_lo  , j_lo+1, k_lo  ]
                        + coef_221 * lo[i_lo+1, j_lo+1, k_lo  ]
                        + coef_112 * lo[i_lo  , j_lo  , k_lo+1]
                        + coef_212 * lo[i_lo+1, j_lo  , k_lo+1]
                        + coef_122 * lo[i_lo  , j_lo+1, k_lo+1]
                        + coef_222 * lo[i_lo+1, j_lo+1, k_lo+1] )
                end
            end
        end
    end
end

function upsample_lin_ref!(hi, lo, x_hi, y_hi, z_hi, x_lo, y_lo, z_lo, is_hi, js_hi, ks_hi, ie_hi, je_hi, ke_hi)
    interp = interpolate((x_lo, y_lo, z_lo), lo, (Gridded(Linear()), Gridded(Linear()), Gridded(Linear())))
    hi[is_hi:ie_hi, js_hi:je_hi, ks_hi:ke_hi] .= interp(x_hi[is_hi:ie_hi], y_hi[js_hi:je_hi], z_hi[ks_hi:ke_hi])
end

## Set up the grids.
itot_lo = 256; jtot_lo = 192; ktot_lo = 128
itot_hi = 512; jtot_hi = 384; ktot_hi = 256
igc_lo = 1; jgc_lo = 1; kgc_lo = 1
igc_hi = 1; jgc_hi = 1; kgc_hi = 1

# Derive dimensions
ifac = itot_hi÷itot_lo; jfac = jtot_hi÷jtot_lo; kfac = ktot_hi÷ktot_lo;
icells_lo = itot_lo + 2igc_lo; jcells_lo = jtot_lo + 2jgc_lo; kcells_lo = ktot_lo + 2kgc_lo
icells_hi = itot_hi + 2igc_hi; jcells_hi = jtot_hi + 2jgc_hi; kcells_hi = ktot_hi + 2kgc_hi
is_lo = igc_lo+1; js_lo = jgc_lo+1; ks_lo = kgc_lo+1
is_hi = igc_hi+1; js_hi = jgc_hi+1; ks_hi = kgc_hi+1
ie_lo = igc_lo+itot_lo; je_lo = jgc_lo+jtot_lo; ke_lo = kgc_lo+ktot_lo
ie_hi = igc_hi+itot_hi; je_hi = jgc_hi+jtot_hi; ke_hi = kgc_hi+ktot_hi

a_lo = rand(icells_lo, jcells_lo, kcells_lo)
a_hi_ref = zeros(icells_hi, jcells_hi, kcells_hi)
a_hi = zeros(icells_hi, jcells_hi, kcells_hi)
a_hi_fast = zeros(icells_hi, jcells_hi, kcells_hi)

dx_lo = 1/itot_lo; dy_lo = 1/jtot_lo; dz_lo = 1/ktot_lo
x_lo = collect(range(0.5-igc_lo, length=icells_lo)) .* dx_lo
y_lo = collect(range(0.5-jgc_lo, length=jcells_lo)) .* dy_lo
z_lo = collect(range(0.5-kgc_lo, length=kcells_lo)) .* dz_lo
xh_lo = collect(range(-igc_lo, length=icells_lo)) .* dx_lo
yh_lo = collect(range(-jgc_lo, length=jcells_lo)) .* dy_lo

dx_hi = 1/itot_hi; dy_hi = 1/jtot_hi; dz_hi = 1/ktot_hi
x_hi = collect(range(0.5-igc_hi, length=icells_hi)) .* dx_hi
y_hi = collect(range(0.5-jgc_hi, length=jcells_hi)) .* dy_hi
z_hi = collect(range(0.5-kgc_hi, length=kcells_hi)) .* dz_hi
xh_hi = collect(range(-igc_hi, length=icells_hi)) .* dx_hi
yh_hi = collect(range(-jgc_hi, length=jcells_hi)) .* dy_hi

## Compute a reference using Interpolations.jl
@btime upsample_lin_ref!(
    a_hi_ref, a_lo, x_hi, y_hi, z_hi, x_lo, y_lo, z_lo,
    is_hi, js_hi, ks_hi, ie_hi, je_hi, ke_hi)
@btime upsample_lin!(a_hi, a_lo, itot_lo, jtot_lo, ktot_lo,
    is_hi, js_hi, ks_hi, is_lo, js_lo, ks_lo,
    ifac, jfac, kfac, 0, 0, 0)
@upsample_lin_fast("222", 2, 2, 2, 0, 0, 0)
@btime upsample_lin_222!(
    a_hi_fast, a_lo, itot_lo, jtot_lo, ktot_lo,
    is_hi, js_hi, ks_hi, is_lo, js_lo, ks_lo)

a_lo_int = @view a_lo[is_lo:ie_lo, js_lo:je_lo, ks_lo:ke_lo]
a_hi_ref_int = @view a_hi_ref[is_hi:ie_hi, js_hi:je_hi, ks_hi:ke_hi]
a_hi_int = @view a_hi[is_hi:ie_hi, js_hi:je_hi, ks_hi:ke_hi]
a_hi_fast_int = @view a_hi_fast[is_hi:ie_hi, js_hi:je_hi, ks_hi:ke_hi]

println("Standard equal to ref: ", a_hi_int ≈ a_hi_ref_int)
println("Fast equal to ref: ", a_hi_fast_int ≈ a_hi_ref_int)

## Plot the output.
x_lo_int = @view x_lo[is_lo:ie_lo]
x_hi_int = @view x_hi[is_hi:ie_hi]

figure()
plot(x_hi_int, a_hi_int[:, 1, 1], "C0-o")
plot(x_hi_int, a_hi_ref_int[:, 1, 1], "C1-^")
plot(x_hi_int, a_hi_fast_int[:, 1, 1], "C2-*")
plot(x_lo_int, a_lo_int[:, 1, 1], "k:+")
tight_layout()
display(gcf())

## Compute a reference using Interpolations.jl
u_lo = rand(icells_lo, jcells_lo, kcells_lo)
u_hi_ref = zeros(icells_hi, jcells_hi, kcells_hi)
u_hi = zeros(icells_hi, jcells_hi, kcells_hi)
u_hi_fast = zeros(icells_hi, jcells_hi, kcells_hi)

@btime upsample_lin_ref!(
     u_hi_ref, u_lo, xh_hi, y_hi, z_hi, xh_lo, y_lo, z_lo,
     is_hi, js_hi, ks_hi, ie_hi, je_hi, ke_hi)
@btime upsample_lin!(
    u_hi, u_lo, itot_lo, jtot_lo, ktot_lo,
    is_hi, js_hi, ks_hi, is_lo, js_lo, ks_lo,
    ifac, jfac, kfac, 0.5, 0, 0)
@upsample_lin_fast("222_u", 2, 2, 2, 0.5, 0, 0)
@btime upsample_lin_222_u!(
    u_hi_fast, u_lo, itot_lo, jtot_lo, ktot_lo,
    is_hi, js_hi, ks_hi, is_lo, js_lo, ks_lo)

u_lo_int = @view a_lo[is_lo:ie_lo, js_lo:je_lo, ks_lo:ke_lo]
u_hi_ref_int = @view a_hi_ref[is_hi:ie_hi, js_hi:je_hi, ks_hi:ke_hi]
u_hi_int = @view a_hi[is_hi:ie_hi, js_hi:je_hi, ks_hi:ke_hi]
u_hi_fast_int = @view a_hi_fast[is_hi:ie_hi, js_hi:je_hi, ks_hi:ke_hi]

println("Standard equal to ref: ", u_hi_int ≈ u_hi_ref_int)
println("Fast equal to ref: ", u_hi_fast_int ≈ u_hi_ref_int)

## Plot the output.
xh_lo_int = @view xh_lo[is_lo:ie_lo]
xh_hi_int = @view xh_hi[is_hi:ie_hi]

figure()
plot(xh_hi_int, u_hi_int[:, 1, 1], "C0-o")
plot(xh_hi_int, u_hi_ref_int[:, 1, 1], "C1-^")
plot(xh_hi_int, u_hi_fast_int[:, 1, 1], "C2-*")
plot(xh_lo_int, u_lo_int[:, 1, 1], "k:+")
tight_layout()
display(gcf())

show()


text:
function fib15()
#
#  test polynomial suggested by Goedecker
#
    n = 15;
    p = fib(n)
    z = [-.9115849521765970-.1896359497733640*im;
         -.9115849521765970+.1896359497733640*im;
         -.7620785395636261-.5388639564255547*im;
         -.7620785395636261+.5388639564255547*im;
         -.4858038058618722-.8023902123527991*im;
         -.4858038058618722+.8023902123527991*im;
         -.1248130033927766-.9371029958793902*im;
         -.1248130033927766+.9371029958793902*im;
         .2657642744875577-.9183694775911604*im;
         .2657642744875577+.9183694775911604*im;
         .6253568882929779-.7424114494678755*im;
         .6253568882929779+.7424114494678755*im;
         .8931744004970845-.4240640341420445*im;
         .8931744004970845+.4240640341420445*im;
         1.999969475434503];

    z = [z ones(n)];
    p, PolyZeros(z)
end


text:
using Base.CoreLogging: LogLevel, AbstractLogger,
    BelowMinLevel, Debug, Info, Warn, Error, AboveMaxLevel

import Base.CoreLogging: handle_message, shouldlog, min_enabled_level, catch_exceptions

using Logging: default_metafmt, termlength

"""
    OhMyLogger(stream=stderr, min_level=Info; meta_formatter=default_metafmt,
                  show_limited=true, right_justify=0)

Logger with formatting optimized for readability in a text console, for example
interactive work with the Julia REPL.

Log levels less than `min_level` are filtered out.

Message formatting can be controlled by setting keyword arguments:

* `meta_formatter` is a function which takes the log event metadata
  `(level, _module, group, id, file, line)` and returns a color (as would be
  passed to printstyled), prefix and suffix for the log message.  The
  default is to prefix with the log level and a suffix containing the module,
  file and line location.
* `show_limited` limits the printing of large data structures to something
  which can fit on the screen by setting the `:limit` `IOContext` key during
  formatting.
* `right_justify` is the integer column which log metadata is right justified
  at. The default is zero (metadata goes on its own line).
"""
struct OhMyLogger <: AbstractLogger
    stream::IO
    min_level::LogLevel
    meta_formatter
    show_limited::Bool
    right_justify::Int
    message_limits::Dict{Any,Int}

    last_message_nlines::Ref{Int}
    last_message_id::Ref{Symbol}
end

function OhMyLogger(stream::IO=stderr, min_level=Info;
                       meta_formatter=default_metafmt, show_limited=true,
                       right_justify=0)

    OhMyLogger(stream, min_level, meta_formatter,
               show_limited, right_justify, Dict{Any,Int}(),
               Ref(0), Ref(Symbol(""))
              )
end

shouldlog(logger::OhMyLogger, level, _module, group, id) =
    get(logger.message_limits, id, 1) > 0

min_enabled_level(logger::OhMyLogger) = logger.min_level

##################### Formatting of values in key value pairs #######################
showvalue(io, msg, key::Symbol) = showvalue(io, msg, Val{key}())
_showvalue(io, msg) = show(io, "text/plain", msg)

function showvalue(io, msg, ::Any)
    if msg isa AbstractString && occursin(r"\%\s*$", msg) # ends with a %
        showvalue(io, msg, Val{:progress}())
    else
        _showvalue(io, msg)
    end
end

showvalue(io, msg, ::Val{:progress}) = _showvalue(io, progress_message(msg))

function showvalue(io, e::Tuple{Exception,Any}, ::Any)
    ex,bt = e
    showerror(io, ex, bt; backtrace = bt!=nothing)
end
showvalue(io, ex::Exception, ::Any) = showerror(io, ex)

#####################################################################################

function prepare_output(logger::OhMyLogger, level, message, _module, group, id,
                        filepath, line; kwargs...)
    # Generate a text representation of the message and all key value pairs,
    # split into lines.
    msglines = [(indent=0,msg=l) for l in split(chomp(string(message)), '\n')]
    dsize = displaysize(logger.stream)
    if !isempty(kwargs)
        valbuf = IOBuffer()
        rows_per_value = max(1, dsize[1]÷(length(kwargs)+1))
        valio = IOContext(IOContext(valbuf, logger.stream),
                          :displaysize=>(rows_per_value,dsize[2]-5))
        if logger.show_limited
            valio = IOContext(valio, :limit=>true)
        end
        for (key,val) in pairs(kwargs)
            showvalue(valio, val, key)
            vallines = split(String(take!(valbuf)), '\n')
            if length(vallines) == 1
                push!(msglines, (indent=2,msg=SubString("$key = $(vallines[1])")))
            else
                push!(msglines, (indent=2,msg=SubString("$key =")))
                append!(msglines, ((indent=3,msg=line) for line in vallines))
            end
        end
    end

    # Format lines as text with appropriate indentation and with a box
    # decoration on the left.
    color, prefix, suffix = logger.meta_formatter(level, _module, group, id, filepath, line)
    minsuffixpad = 2
    buf = IOBuffer()
    iob = IOContext(buf, logger.stream)
    nonpadwidth = 2 + (isempty(prefix) || length(msglines) > 1 ? 0 : length(prefix)+1) +
                  msglines[end].indent + termlength(msglines[end].msg) +
                  (isempty(suffix) ? 0 : length(suffix)+minsuffixpad)
    justify_width = min(logger.right_justify, dsize[2])
    if nonpadwidth > justify_width && !isempty(suffix)
        push!(msglines, (indent=0,msg=SubString("")))
        minsuffixpad = 0
        nonpadwidth = 2 + length(suffix)
    end
    for (i,(indent,msg)) in enumerate(msglines)
        boxstr = length(msglines) == 1 ? "[ " :
                 i == 1                ? "┌ " :
                 i < length(msglines)  ? "│ " :
                                         "└ "
        printstyled(iob, boxstr, bold=true, color=color)
        if i == 1 && !isempty(prefix)
            printstyled(iob, prefix, " ", bold=true, color=color)
        end
        print(iob, " "^indent, msg)
        if i == length(msglines) && !isempty(suffix)
            npad = max(0, justify_width - nonpadwidth) + minsuffixpad
            print(iob, " "^npad)
            printstyled(iob, suffix, color=Base.stackframe_lineinfo_color())
        end
        println(iob)
    end

    buf
end

function handle_message(logger::OhMyLogger, level, message, _module, group, id,
                        filepath, line; maxlog=nothing, overwrite_lastlog=nothing,
                        kwargs...)
    id = Symbol(filepath, "#", line)  # HACK around https://github.com/JuliaLang/julia/issues/29227
    prev_id = logger.last_message_id[]
    if overwrite_lastlog == nothing
        # default to only overwrite if same source
        overwrite_lastlog = id == prev_id
    end

    if maxlog != nothing && maxlog isa Integer
        if overwrite_lastlog && haskey(logger.message_limits, prev_id)
            # We are overwriting it so it undoes one of it's limitted uses
            logger.message_limits[prev_id] += 1
        end

        remaining = get!(logger.message_limits, id, maxlog)
        logger.message_limits[id] = remaining - 1
        remaining > 0 || return
    end

    buf = prepare_output(logger, level, message, _module, group, id,
                         filepath, line; kwargs...)

    buf_contents = take!(buf)
    if overwrite_lastlog
        move_cursor_up_while_clearing_lines(logger.stream, logger.last_message_nlines[])
    end

    write(logger.stream, buf_contents)
    logger.last_message_id[] = id
    logger.last_message_nlines[] = count(isequal(Int('\n')), buf_contents)
    nothing
end


text:
# Extend Zygote to work with static vectors and custom types on the GPU
# Here be dragons

using ForwardDiff: Chunk, Dual, dualize, partials, value
using Zygote: unbroadcast

Zygote.accum(x::AbstractArray{<:SizedVector}, ys::AbstractArray{<:SVector}...) = Zygote.accum.(convert(typeof(ys[1]), x), ys...)
Zygote.accum(x::AbstractArray{<:SVector}, ys::AbstractArray{<:SizedVector}...) = Zygote.accum.(x, convert.(typeof(x), ys)...)

Zygote.accum(x::Vector{<:SVector} , y::CuArray{<:SVector}) = Zygote.accum(cu(x), y)
Zygote.accum(x::CuArray{<:SVector}, y::Vector{<:SVector} ) = Zygote.accum(x, cu(y))

Base.:+(x::Real, y::SizedVector) = x .+ y
Base.:+(x::SizedVector, y::Real) = x .+ y

Base.:+(x::Real, y::Zygote.OneElement) = x .+ y
Base.:+(x::Zygote.OneElement, y::Real) = x .+ y

function Base.:+(x::Atom{T, T, T, T}, y::Atom{T, T, T, T}) where T
    Atom{T, T, T, T}(0, x.charge + y.charge, x.mass + y.mass, x.σ + y.σ, x.ϵ + y.ϵ, false)
end

function Base.:-(x::Atom{T, T, T, T}, y::Atom{T, T, T, T}) where T
    Atom{T, T, T, T}(0, x.charge - y.charge, x.mass - y.mass, x.σ - y.σ, x.ϵ - y.ϵ, false)
end

function Zygote.accum(x::LennardJones{S, C, W, WS, F, E}, y::LennardJones{S, C, W, WS, F, E}) where {S, C, W, WS, F, E}
    LennardJones{S, C, W, WS, F, E}(
        x.cutoff,
        x.nl_only,
        x.lorentz_mixing,
        x.weight_14 + y.weight_14,
        x.weight_solute_solvent + y.weight_solute_solvent,
        x.force_units,
        x.energy_units,
    )
end

function Zygote.accum(x::CoulombReactionField{D, S, W, T, F, E}, y::CoulombReactionField{D, S, W, T, F, E}) where {D, S, W, T, F, E}
    CoulombReactionField{D, S, W, T, F, E}(x.dist_cutoff + y.dist_cutoff, x.solvent_dielectric + y.solvent_dielectric, x.nl_only,
                x.weight_14 + y.weight_14, x.coulomb_const + y.coulomb_const, x.force_units, x.energy_units)
end

function Zygote.accum_sum(xs::AbstractArray{Tuple{LennardJones{S, C, W, WS, F, E}, CoulombReactionField{D, SO, WC, T, F, E}}}; dims=:) where {S, C, W, WS, F, E, D, SO, WC, T}
    reduce(Zygote.accum, xs, dims=dims; init=(
        LennardJones{S, C, W, WS, F, E}(nothing, false, false, zero(W), zero(WS), NoUnits, NoUnits),
        CoulombReactionField{D, SO, WC, T, F, E}(zero(D), zero(SO), false, zero(WC), zero(T), NoUnits, NoUnits),
    ))
end

function Zygote.accum(x::Tuple{NTuple{N, Int}, NTuple{N, T}, NTuple{N, E}, Bool},
                        y::Tuple{NTuple{N, Int}, NTuple{N, T}, NTuple{N, E}, Bool}) where {N, T, E}
    ntuple(n -> 0, N), x[2] .+ y[2], x[3] .+ y[3], false
end

Base.zero(::Type{Atom{T, T, T, T}}) where {T} = Atom(0, zero(T), zero(T), zero(T), zero(T), false)
atom_or_empty(at::Atom, T) = at
atom_or_empty(at::Nothing, T) = zero(Atom{T, T, T, T})

Zygote.z2d(dx::AbstractArray{Union{Nothing, Atom{T, T, T, T}}}, primal::AbstractArray{Atom{T, T, T, T}}) where {T} = atom_or_empty.(dx, T)
Zygote.z2d(dx::SVector{3, T}, primal::T) where {T} = sum(dx)

Zygote.unbroadcast(x::Tuple{Any}, x̄::Nothing) = nothing

function Zygote.unbroadcast(x::AbstractArray{<:Real}, x̄::AbstractArray{<:StaticVector})
    if length(x) == length(x̄)
        Zygote._project(x, sum.(x̄))
    else
        dims = ntuple(d -> size(x, d) == 1 ? d : ndims(x̄) + 1, ndims(x̄))
        Zygote._project(x, accum_sum(x̄; dims=dims))
    end
end

Zygote._zero(xs::AbstractArray{<:StaticVector}, T) = fill!(similar(xs, T), zero(T))

function Zygote._zero(xs::AbstractArray{Atom{T, T, T, T}}, ::Type{Atom{T, T, T, T}}) where {T}
    fill!(similar(xs), Atom{T, T, T, T}(0, zero(T), zero(T), zero(T), zero(T), false))
end

function Base.zero(::Type{Union{Nothing, SizedVector{D, T, Vector{T}}}}) where {D, T}
    zero(SizedVector{D, T, Vector{T}})
end

function Zygote._backmean(xs::AbstractArray{SVector{D, T}}, Δ::SVector{D, T}, ::Colon) where {D, T}
    zero(xs) .+ (Δ ./ length(xs),)
end

# Slower version than in Zygote but doesn't give wrong gradients on the GPU for repeated indices
# Here we just move it to the CPU then move it back
# See https://github.com/FluxML/Zygote.jl/pull/1131
Zygote.∇getindex(x::CuArray, inds::Tuple{AbstractArray{<:Integer}}) = dy -> begin
    inds1_cpu = Array(inds[1])
    dx = zeros(eltype(dy), length(x))
    dxv = view(dx, inds1_cpu)
    dxv .= Zygote.accum.(dxv, Zygote._droplike(Array(dy), dxv))
    return Zygote._project(x, cu(dx)), nothing
end

# Extend to add extra empty partials before (B) and after (A) the SVector partials
@generated function ForwardDiff.dualize(::Type{T}, x::StaticArray, ::Val{B}, ::Val{A}) where {T, B, A}
    N = length(x)
    dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i + $B}())) for i in 1:N]...)
    V = StaticArrays.similar_type(x, Dual{T, eltype(x), N + B + A})
    return quote
        chunk = Chunk{$N + $B + $A}()
        $(Expr(:meta, :inline))
        return $V($(dx))
    end
end

# Dualize a value with extra partials
macro dualize(x, n_partials::Integer, active_partial::Integer)
    ps = [i == active_partial for i in 1:n_partials]
    return :(ForwardDiff.Dual($(esc(x)), $(ps...)))
end

# No gradient for cutoff type
function dualize_fb(inter::LennardJones{S, C, W, WS, F, E}) where {S, C, W, WS, F, E}
    w14, wss = inter.weight_14, inter.weight_solute_solvent
    dual_weight_14 = @dualize(w14, 23, 1)
    dual_weight_solute_solvent = @dualize(wss, 23, 2)
    return LennardJones{S, C, typeof(dual_weight_14), typeof(dual_weight_solute_solvent), F, E}(
                inter.cutoff, inter.nl_only, inter.lorentz_mixing, dual_weight_14,
                dual_weight_solute_solvent, inter.force_units, inter.energy_units)
end

function dualize_fb(inter::CoulombReactionField{D, S, W, T, F, E}) where {D, S, W, T, F, E}
    dc, sd, w14, cc = inter.dist_cutoff, inter.solvent_dielectric, inter.weight_14, inter.coulomb_const
    dual_dist_cutoff        = @dualize(dc , 23, 3)
    dual_solvent_dielectric = @dualize(sd , 23, 4)
    dual_weight_14          = @dualize(w14, 23, 5)
    dual_coulomb_const      = @dualize(cc , 23, 6)
    return CoulombReactionField{typeof(dual_dist_cutoff), typeof(dual_solvent_dielectric), typeof(dual_weight_14), typeof(dual_coulomb_const), F, E}(
                                dual_dist_cutoff, dual_solvent_dielectric, inter.nl_only, dual_weight_14,
                                dual_coulomb_const, inter.force_units, inter.energy_units)
end

function dualize_fb(inter::HarmonicBond{D, K}) where {D, K}
    b0, kb = inter.b0, inter.kb
    dual_b0 = @dualize(b0, 11, 1)
    dual_kb = @dualize(kb, 11, 2)
    return HarmonicBond{typeof(dual_b0), typeof(dual_kb)}(dual_b0, dual_kb)
end

function dualize_fb(inter::HarmonicAngle{D, K}) where {D, K}
    th0, cth = inter.th0, inter.cth
    dual_th0 = @dualize(th0, 14, 1)
    dual_cth = @dualize(cth, 14, 2)
    return HarmonicAngle{typeof(dual_th0), typeof(dual_cth)}(dual_th0, dual_cth)
end

function dualize_fb(inter::PeriodicTorsion{6, T, E}) where {T, E}
    p1, p2, p3, p4, p5, p6 = inter.phases
    k1, k2, k3, k4, k5, k6 = inter.ks
    dual_phases = (
        @dualize(p1, 27,  1), @dualize(p2, 27,  2), @dualize(p3, 27,  3),
        @dualize(p4, 27,  4), @dualize(p5, 27,  5), @dualize(p6, 27,  6),
    )
    dual_ks = (
        @dualize(k1, 27,  7), @dualize(k2, 27,  8), @dualize(k3, 27,  9),
        @dualize(k4, 27, 10), @dualize(k5, 27, 11), @dualize(k6, 27, 12),
    )
    return PeriodicTorsion{6, eltype(dual_phases), eltype(dual_ks)}(inter.periodicities,
                            dual_phases, dual_ks, inter.proper)
end

function dualize_atom_fb1(at::Atom)
    c, m, σ, ϵ = at.charge, at.mass, at.σ, at.ϵ
    dual_charge = @dualize(c, 23, 13)
    dual_mass = @dualize(m, 23, 14)
    dual_σ = @dualize(σ, 23, 15)
    dual_ϵ = @dualize(ϵ, 23, 16)
    return Atom{typeof(dual_charge), typeof(dual_mass), typeof(dual_σ), typeof(dual_ϵ)}(
                at.index, dual_charge, dual_mass, dual_σ, dual_ϵ, at.solute)
end

function dualize_atom_fb2(at::Atom)
    c, m, σ, ϵ = at.charge, at.mass, at.σ, at.ϵ
    dual_charge = @dualize(c, 23, 17)
    dual_mass = @dualize(m, 23, 18)
    dual_σ = @dualize(σ, 23, 19)
    dual_ϵ = @dualize(ϵ, 23, 20)
    return Atom{typeof(dual_charge), typeof(dual_mass), typeof(dual_σ), typeof(dual_ϵ)}(
                at.index, dual_charge, dual_mass, dual_σ, dual_ϵ, at.solute)
end

function dual_function_svec(f::F) where F
    function (arg1)
        ds1 = dualize(Nothing, arg1, Val(0), Val(0))
        return f(ds1)
    end
end

function dual_function_svec_real(f::F) where F
    function (arg1::SVector{D, T}, arg2) where {D, T}
        ds1 = dualize(Nothing, arg1, Val(0), Val(1))
        # Leaving the integer type in here results in Float32 -> Float64 conversion
        ds2 = Zygote.dual(arg2 isa Int ? T(arg2) : arg2, (false, false, false, true))
        return f(ds1, ds2)
    end
end

function dual_function_svec_svec(f::F) where F
    function (arg1, arg2)
        ds1 = dualize(Nothing, arg1, Val(0), Val(3))
        ds2 = dualize(Nothing, arg2, Val(3), Val(0))
        return f(ds1, ds2)
    end
end

function dual_function_atom(f::F) where F
    function (arg1)
        c, m, σ, ϵ = arg1.charge, arg1.mass, arg1.σ, arg1.ϵ
        ds1 = Atom(arg1.index, @dualize(c, 4, 1), @dualize(m, 4, 2), @dualize(σ, 4, 3),
                    @dualize(ϵ, 4, 4), arg1.solute)
        return f(ds1)
    end
end

function dual_function_force_broadcast(f::F) where F
    function (arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8)
        ds1 = map(dualize_fb, arg1)
        ds2 = dualize(Nothing, arg2, Val(6), Val(14))
        ds3 = dualize(Nothing, arg3, Val(9), Val(11))
        ds4 = dualize_atom_fb1(arg4)
        ds5 = dualize_atom_fb2(arg5)
        ds6 = dualize(Nothing, arg6, Val(20), Val(0))
        ds7 = arg7
        ds8 = arg8
        return f(ds1, ds2, ds3, ds4, ds5, ds6, ds7, ds8)
    end
end

function dual_function_specific_2_atoms(f::F) where F
    function (arg1, arg2, arg3, arg4)
        ds1 = dualize_fb(arg1)
        ds2 = dualize(Nothing, arg2, Val(2), Val(6))
        ds3 = dualize(Nothing, arg3, Val(5), Val(3))
        ds4 = dualize(Nothing, arg4, Val(8), Val(0))
        return f(ds1, ds2, ds3, ds4)
    end
end

function dual_function_specific_3_atoms(f::F) where F
    function (arg1, arg2, arg3, arg4, arg5)
        ds1 = dualize_fb(arg1)
        ds2 = dualize(Nothing, arg2, Val( 2), Val(9))
        ds3 = dualize(Nothing, arg3, Val( 5), Val(6))
        ds4 = dualize(Nothing, arg4, Val( 8), Val(3))
        ds5 = dualize(Nothing, arg5, Val(11), Val(0))
        return f(ds1, ds2, ds3, ds4, ds5)
    end
end

function dual_function_specific_4_atoms(f::F) where F
    function (arg1, arg2, arg3, arg4, arg5, arg6)
        ds1 = dualize_fb(arg1)
        ds2 = dualize(Nothing, arg2, Val(12), Val(12))
        ds3 = dualize(Nothing, arg3, Val(15), Val( 9))
        ds4 = dualize(Nothing, arg4, Val(18), Val( 6))
        ds5 = dualize(Nothing, arg5, Val(21), Val( 3))
        ds6 = dualize(Nothing, arg6, Val(24), Val( 0))
        return f(ds1, ds2, ds3, ds4, ds5, ds6)
    end
end

function dual_function_born_radii_loop(f::F) where F
    function (arg1, arg2, arg3, arg4, arg5, arg6)
        ds1 = dualize(Nothing, arg1, Val(0), Val(8))
        ds2 = dualize(Nothing, arg2, Val(3), Val(5))
        ds3 = Zygote.dual(arg3, (false, false, false, false, false, false, true , false, false, false, false))
        ds4 = Zygote.dual(arg4, (false, false, false, false, false, false, false, true , false, false, false))
        ds5 = arg5
        ds6 = dualize(Nothing, arg6, Val(8), Val(0))
        return f(ds1, ds2, ds3, ds4, ds5, ds6)
    end
end

function dual_function_gb_force_loop_1(f::F) where F
    function (arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11)
        ds1  = dualize(Nothing, arg1, Val(0), Val(11))
        ds2  = dualize(Nothing, arg2, Val(3), Val( 8))
        ds3  = arg3
        ds4  = arg4
        ds5  = Zygote.dual(arg5 , (false, false, false, false, false, false, true , false, false, false, false, false, false, false))
        ds6  = Zygote.dual(arg6 , (false, false, false, false, false, false, false, true , false, false, false, false, false, false))
        ds7  = Zygote.dual(arg7 , (false, false, false, false, false, false, false, false, true , false, false, false, false, false))
        ds8  = Zygote.dual(arg8 , (false, false, false, false, false, false, false, false, false, true , false, false, false, false))
        ds9  = arg9
        ds10 = Zygote.dual(arg10, (false, false, false, false, false, false, false, false, false, false, true , false, false, false))
        ds11 = dualize(Nothing, arg11, Val(11), Val(0))
        return f(ds1, ds2, ds3, ds4, ds5, ds6, ds7, ds8, ds9, ds10, ds11)
    end
end

function dual_function_gb_force_loop_2(f::F) where F
    function (arg1, arg2, arg3, arg4, arg5, arg6, arg7)
        ds1 = dualize(Nothing, arg1, Val(0), Val(9))
        ds2 = dualize(Nothing, arg2, Val(3), Val(6))
        ds3 = Zygote.dual(arg3, (false, false, false, false, false, false, true , false, false, false, false, false))
        ds4 = Zygote.dual(arg4, (false, false, false, false, false, false, false, true , false, false, false, false))
        ds5 = Zygote.dual(arg5, (false, false, false, false, false, false, false, false, true , false, false, false))
        ds6 = arg6
        ds7 = dualize(Nothing, arg7, Val(9), Val(0))
        return f(ds1, ds2, ds3, ds4, ds5, ds6, ds7)
    end
end

function dual_function_gb_energy_loop(f::F) where F
    function (arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10,
                arg11, arg12, arg13, arg14, arg15, arg16)
        ds1  = dualize(Nothing, arg1, Val(0), Val(15))
        ds2  = dualize(Nothing, arg2, Val(3), Val(12))
        ds3  = arg3
        ds4  = arg4
        ds5  = Zygote.dual(arg5 , (false, false, false, false, false, false, true , false, false, false, false, false, false, false, false, false, false, false))
        ds6  = Zygote.dual(arg6 , (false, false, false, false, false, false, false, true , false, false, false, false, false, false, false, false, false, false))
        ds7  = Zygote.dual(arg7 , (false, false, false, false, false, false, false, false, true , false, false, false, false, false, false, false, false, false))
        ds8  = Zygote.dual(arg8 , (false, false, false, false, false, false, false, false, false, true , false, false, false, false, false, false, false, false))
        ds9  = Zygote.dual(arg9 , (false, false, false, false, false, false, false, false, false, true , true , false, false, false, false, false, false, false))
        ds10 = arg10
        ds11 = Zygote.dual(arg11, (false, false, false, false, false, false, false, false, false, false, false, true , false, false, false, false, false, false))
        ds12 = Zygote.dual(arg12, (false, false, false, false, false, false, false, false, false, false, false, false, true , false, false, false, false, false))
        ds13 = Zygote.dual(arg13, (false, false, false, false, false, false, false, false, false, false, false, false, false, true , false, false, false, false))
        ds14 = Zygote.dual(arg14, (false, false, false, false, false, false, false, false, false, false, false, false, false, false, true , false, false, false))
        ds15 = arg15
        ds16 = dualize(Nothing, arg16, Val(15), Val(0))
        return f(ds1, ds2, ds3, ds4, ds5, ds6, ds7, ds8, ds9, ds10,
                    ds11, ds12, ds13, ds14, ds15, ds16)
    end
end

@inline function sum_partials(sv::SVector{3, Dual{Nothing, T, P}}, y1, i::Integer) where {T, P}
    partials(sv[1], i) * y1[1] + partials(sv[2], i) * y1[2] + partials(sv[3], i) * y1[3]
end

sized_to_static(v::SizedVector{3, T, Vector{T}}) where {T} = SVector{3, T}(v[1], v[2], v[3])
sized_to_static(v::SizedVector{2, T, Vector{T}}) where {T} = SVector{2, T}(v[1], v[2])

function modify_grad(ȳ_in::AbstractArray{SizedVector{D, T, Vector{T}}}, arg::CuArray) where {D, T}
    cu(sized_to_static.(ȳ_in))
end

function modify_grad(ȳ_in::AbstractArray{SizedVector{D, T, Vector{T}}}, arg) where {D, T}
    sized_to_static.(ȳ_in)
end

modify_grad(ȳ_in, arg::CuArray) = cu(ȳ_in)
modify_grad(ȳ_in, arg) = ȳ_in

@inline function Zygote.broadcast_forward(f, arg1::AbstractArray{SVector{D, T}}) where {D, T}
    out = dual_function_svec(f).(arg1)
    y = map(x -> value.(x), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        barg1 = broadcast(ȳ, out) do y1, o1
            if length(y1) == 1
                y1 .* SVector{D, T}(partials(o1))
            else
                SVector{D, T}(sum_partials(o1, y1, 1), sum_partials(o1, y1, 2), sum_partials(o1, y1, 3))
            end
        end
        darg1 = unbroadcast(arg1, barg1)
        (nothing, nothing, darg1)
    end
    return y, bc_fwd_back
end

@inline function Zygote.broadcast_forward(f, arg1::AbstractArray{SVector{D, T}}, arg2) where {D, T}
    out = dual_function_svec_real(f).(arg1, arg2)
    y = map(x -> value.(x), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        barg1 = broadcast(ȳ, out) do y1, o1
            if length(y1) == 1
                y1 .* SVector{D, T}(partials.((o1,), (1, 2, 3)))
            else
                SVector{D, T}(sum_partials(o1, y1, 1), sum_partials(o1, y1, 2), sum_partials(o1, y1, 3))
            end
        end
        darg1 = unbroadcast(arg1, barg1)
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> y1 .* partials.(o1, 4), ȳ, out))
        (nothing, nothing, darg1, darg2)
    end
    return y, bc_fwd_back
end

@inline function Zygote.broadcast_forward(f, arg1::AbstractArray{SVector{D, T}}, arg2::AbstractArray{SVector{D, T}}) where {D, T}
    out = dual_function_svec_svec(f).(arg1, arg2)
    y = map(x -> value.(x), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        barg1 = broadcast(ȳ, out) do y1, o1
            if length(y1) == 1
                y1 .* SVector{D, T}(partials.((o1,), (1, 2, 3)))
            else
                SVector{D, T}(sum_partials(o1, y1, 1), sum_partials(o1, y1, 2), sum_partials(o1, y1, 3))
            end
        end
        darg1 = unbroadcast(arg1, barg1)
        barg2 = broadcast(ȳ, out) do y1, o1
            if length(y1) == 1
                y1 .* SVector{D, T}(partials.((o1,), (4, 5, 6)))
            else
                SVector{D, T}(sum_partials(o1, y1, 4), sum_partials(o1, y1, 5), sum_partials(o1, y1, 6))
            end
        end
        darg2 = unbroadcast(arg2, barg2)
        (nothing, nothing, darg1, darg2)
    end
    return y, bc_fwd_back
end

# For mass, charge etc.
@inline function Zygote.broadcast_forward(f, arg1::AbstractArray{<:Atom})
    out = dual_function_atom(f).(arg1)
    y = map(x -> value.(x), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        barg1 = broadcast(ȳ, out) do y1, o1
            ps = partials(o1)
            Atom(0, y1 * ps[1], y1 * ps[2], y1 * ps[3], y1 * ps[4], false)
        end
        darg1 = unbroadcast(arg1, barg1)
        (nothing, nothing, darg1)
    end
    return y, bc_fwd_back
end

function combine_dual_PairwiseInteraction(y1::SVector{3, T}, o1::SVector{3, Dual{Nothing, T, P}}, i::Integer) where {T, P}
    (
        LennardJones{false, Nothing, T, T, typeof(NoUnits), typeof(NoUnits)}(
            nothing, false, false,
            y1[1] * partials(o1[1], i    ) + y1[2] * partials(o1[2], i    ) + y1[3] * partials(o1[3], i    ),
            y1[1] * partials(o1[1], i + 1) + y1[2] * partials(o1[2], i + 1) + y1[3] * partials(o1[3], i + 1),
            NoUnits, NoUnits,
        ),
        CoulombReactionField{T, T, T, T, typeof(NoUnits), typeof(NoUnits)}(
            y1[1] * partials(o1[1], i + 2) + y1[2] * partials(o1[2], i + 2) + y1[3] * partials(o1[3], i + 2),
            y1[1] * partials(o1[1], i + 3) + y1[2] * partials(o1[2], i + 3) + y1[3] * partials(o1[3], i + 3),
            false,
            y1[1] * partials(o1[1], i + 4) + y1[2] * partials(o1[2], i + 4) + y1[3] * partials(o1[3], i + 4),
            y1[1] * partials(o1[1], i + 5) + y1[2] * partials(o1[2], i + 5) + y1[3] * partials(o1[3], i + 5),
            NoUnits, NoUnits,
        ),
    )
end

function combine_dual_SpecificInteraction(inter::HarmonicBond, y1, o1, i::Integer)
    (y1.f1[1] * partials(o1.f1[1], i    ) + y1.f1[2] * partials(o1.f1[2], i    ) + y1.f1[3] * partials(o1.f1[3], i    ) + y1.f2[1] * partials(o1.f2[1], i    ) + y1.f2[2] * partials(o1.f2[2], i    ) + y1.f2[3] * partials(o1.f2[3], i    ),
     y1.f1[1] * partials(o1.f1[1], i + 1) + y1.f1[2] * partials(o1.f1[2], i + 1) + y1.f1[3] * partials(o1.f1[3], i + 1) + y1.f2[1] * partials(o1.f2[1], i + 1) + y1.f2[2] * partials(o1.f2[2], i + 1) + y1.f2[3] * partials(o1.f2[3], i + 1))
end

function combine_dual_SpecificInteraction(inter::HarmonicAngle, y1, o1, i::Integer)
    (y1.f1[1] * partials(o1.f1[1], i    ) + y1.f1[2] * partials(o1.f1[2], i    ) + y1.f1[3] * partials(o1.f1[3], i    ) + y1.f2[1] * partials(o1.f2[1], i    ) + y1.f2[2] * partials(o1.f2[2], i    ) + y1.f2[3] * partials(o1.f2[3], i    ) + y1.f3[1] * partials(o1.f3[1], i    ) + y1.f3[2] * partials(o1.f3[2], i    ) + y1.f3[3] * partials(o1.f3[3], i    ),
     y1.f1[1] * partials(o1.f1[1], i + 1) + y1.f1[2] * partials(o1.f1[2], i + 1) + y1.f1[3] * partials(o1.f1[3], i + 1) + y1.f2[1] * partials(o1.f2[1], i + 1) + y1.f2[2] * partials(o1.f2[2], i + 1) + y1.f2[3] * partials(o1.f2[3], i + 1) + y1.f3[1] * partials(o1.f3[1], i + 1) + y1.f3[2] * partials(o1.f3[2], i + 1) + y1.f3[3] * partials(o1.f3[3], i + 1))
end

function combine_dual_SpecificInteraction(inter::PeriodicTorsion{6}, y1, o1, i::Integer)
    (
        (0, 0, 0, 0, 0, 0),
        (
            y1.f1[1] * partials(o1.f1[1], i     ) + y1.f1[2] * partials(o1.f1[2], i     ) + y1.f1[3] * partials(o1.f1[3], i     ) + y1.f2[1] * partials(o1.f2[1], i     ) + y1.f2[2] * partials(o1.f2[2], i     ) + y1.f2[3] * partials(o1.f2[3], i     ) + y1.f3[1] * partials(o1.f3[1], i     ) + y1.f3[2] * partials(o1.f3[2], i     ) + y1.f3[3] * partials(o1.f3[3], i     ) + y1.f4[1] * partials(o1.f4[1], i     ) + y1.f4[2] * partials(o1.f4[2], i     ) + y1.f4[3] * partials(o1.f4[3], i     ),
            y1.f1[1] * partials(o1.f1[1], i +  1) + y1.f1[2] * partials(o1.f1[2], i +  1) + y1.f1[3] * partials(o1.f1[3], i +  1) + y1.f2[1] * partials(o1.f2[1], i +  1) + y1.f2[2] * partials(o1.f2[2], i +  1) + y1.f2[3] * partials(o1.f2[3], i +  1) + y1.f3[1] * partials(o1.f3[1], i +  1) + y1.f3[2] * partials(o1.f3[2], i +  1) + y1.f3[3] * partials(o1.f3[3], i +  1) + y1.f4[1] * partials(o1.f4[1], i +  1) + y1.f4[2] * partials(o1.f4[2], i +  1) + y1.f4[3] * partials(o1.f4[3], i +  1),
            y1.f1[1] * partials(o1.f1[1], i +  2) + y1.f1[2] * partials(o1.f1[2], i +  2) + y1.f1[3] * partials(o1.f1[3], i +  2) + y1.f2[1] * partials(o1.f2[1], i +  2) + y1.f2[2] * partials(o1.f2[2], i +  2) + y1.f2[3] * partials(o1.f2[3], i +  2) + y1.f3[1] * partials(o1.f3[1], i +  2) + y1.f3[2] * partials(o1.f3[2], i +  2) + y1.f3[3] * partials(o1.f3[3], i +  2) + y1.f4[1] * partials(o1.f4[1], i +  2) + y1.f4[2] * partials(o1.f4[2], i +  2) + y1.f4[3] * partials(o1.f4[3], i +  2),
            y1.f1[1] * partials(o1.f1[1], i +  3) + y1.f1[2] * partials(o1.f1[2], i +  3) + y1.f1[3] * partials(o1.f1[3], i +  3) + y1.f2[1] * partials(o1.f2[1], i +  3) + y1.f2[2] * partials(o1.f2[2], i +  3) + y1.f2[3] * partials(o1.f2[3], i +  3) + y1.f3[1] * partials(o1.f3[1], i +  3) + y1.f3[2] * partials(o1.f3[2], i +  3) + y1.f3[3] * partials(o1.f3[3], i +  3) + y1.f4[1] * partials(o1.f4[1], i +  3) + y1.f4[2] * partials(o1.f4[2], i +  3) + y1.f4[3] * partials(o1.f4[3], i +  3),
            y1.f1[1] * partials(o1.f1[1], i +  4) + y1.f1[2] * partials(o1.f1[2], i +  4) + y1.f1[3] * partials(o1.f1[3], i +  4) + y1.f2[1] * partials(o1.f2[1], i +  4) + y1.f2[2] * partials(o1.f2[2], i +  4) + y1.f2[3] * partials(o1.f2[3], i +  4) + y1.f3[1] * partials(o1.f3[1], i +  4) + y1.f3[2] * partials(o1.f3[2], i +  4) + y1.f3[3] * partials(o1.f3[3], i +  4) + y1.f4[1] * partials(o1.f4[1], i +  4) + y1.f4[2] * partials(o1.f4[2], i +  4) + y1.f4[3] * partials(o1.f4[3], i +  4),
            y1.f1[1] * partials(o1.f1[1], i +  5) + y1.f1[2] * partials(o1.f1[2], i +  5) + y1.f1[3] * partials(o1.f1[3], i +  5) + y1.f2[1] * partials(o1.f2[1], i +  5) + y1.f2[2] * partials(o1.f2[2], i +  5) + y1.f2[3] * partials(o1.f2[3], i +  5) + y1.f3[1] * partials(o1.f3[1], i +  5) + y1.f3[2] * partials(o1.f3[2], i +  5) + y1.f3[3] * partials(o1.f3[3], i +  5) + y1.f4[1] * partials(o1.f4[1], i +  5) + y1.f4[2] * partials(o1.f4[2], i +  5) + y1.f4[3] * partials(o1.f4[3], i +  5),
        ),
        (
            y1.f1[1] * partials(o1.f1[1], i +  6) + y1.f1[2] * partials(o1.f1[2], i +  6) + y1.f1[3] * partials(o1.f1[3], i +  6) + y1.f2[1] * partials(o1.f2[1], i +  6) + y1.f2[2] * partials(o1.f2[2], i +  6) + y1.f2[3] * partials(o1.f2[3], i +  6) + y1.f3[1] * partials(o1.f3[1], i +  6) + y1.f3[2] * partials(o1.f3[2], i +  6) + y1.f3[3] * partials(o1.f3[3], i +  6) + y1.f4[1] * partials(o1.f4[1], i +  6) + y1.f4[2] * partials(o1.f4[2], i +  6) + y1.f4[3] * partials(o1.f4[3], i +  6),
            y1.f1[1] * partials(o1.f1[1], i +  7) + y1.f1[2] * partials(o1.f1[2], i +  7) + y1.f1[3] * partials(o1.f1[3], i +  7) + y1.f2[1] * partials(o1.f2[1], i +  7) + y1.f2[2] * partials(o1.f2[2], i +  7) + y1.f2[3] * partials(o1.f2[3], i +  7) + y1.f3[1] * partials(o1.f3[1], i +  7) + y1.f3[2] * partials(o1.f3[2], i +  7) + y1.f3[3] * partials(o1.f3[3], i +  7) + y1.f4[1] * partials(o1.f4[1], i +  7) + y1.f4[2] * partials(o1.f4[2], i +  7) + y1.f4[3] * partials(o1.f4[3], i +  7),
            y1.f1[1] * partials(o1.f1[1], i +  8) + y1.f1[2] * partials(o1.f1[2], i +  8) + y1.f1[3] * partials(o1.f1[3], i +  8) + y1.f2[1] * partials(o1.f2[1], i +  8) + y1.f2[2] * partials(o1.f2[2], i +  8) + y1.f2[3] * partials(o1.f2[3], i +  8) + y1.f3[1] * partials(o1.f3[1], i +  8) + y1.f3[2] * partials(o1.f3[2], i +  8) + y1.f3[3] * partials(o1.f3[3], i +  8) + y1.f4[1] * partials(o1.f4[1], i +  8) + y1.f4[2] * partials(o1.f4[2], i +  8) + y1.f4[3] * partials(o1.f4[3], i +  8),
            y1.f1[1] * partials(o1.f1[1], i +  9) + y1.f1[2] * partials(o1.f1[2], i +  9) + y1.f1[3] * partials(o1.f1[3], i +  9) + y1.f2[1] * partials(o1.f2[1], i +  9) + y1.f2[2] * partials(o1.f2[2], i +  9) + y1.f2[3] * partials(o1.f2[3], i +  9) + y1.f3[1] * partials(o1.f3[1], i +  9) + y1.f3[2] * partials(o1.f3[2], i +  9) + y1.f3[3] * partials(o1.f3[3], i +  9) + y1.f4[1] * partials(o1.f4[1], i +  9) + y1.f4[2] * partials(o1.f4[2], i +  9) + y1.f4[3] * partials(o1.f4[3], i +  9),
            y1.f1[1] * partials(o1.f1[1], i + 10) + y1.f1[2] * partials(o1.f1[2], i + 10) + y1.f1[3] * partials(o1.f1[3], i + 10) + y1.f2[1] * partials(o1.f2[1], i + 10) + y1.f2[2] * partials(o1.f2[2], i + 10) + y1.f2[3] * partials(o1.f2[3], i + 10) + y1.f3[1] * partials(o1.f3[1], i + 10) + y1.f3[2] * partials(o1.f3[2], i + 10) + y1.f3[3] * partials(o1.f3[3], i + 10) + y1.f4[1] * partials(o1.f4[1], i + 10) + y1.f4[2] * partials(o1.f4[2], i + 10) + y1.f4[3] * partials(o1.f4[3], i + 10),
            y1.f1[1] * partials(o1.f1[1], i + 11) + y1.f1[2] * partials(o1.f1[2], i + 11) + y1.f1[3] * partials(o1.f1[3], i + 11) + y1.f2[1] * partials(o1.f2[1], i + 11) + y1.f2[2] * partials(o1.f2[2], i + 11) + y1.f2[3] * partials(o1.f2[3], i + 11) + y1.f3[1] * partials(o1.f3[1], i + 11) + y1.f3[2] * partials(o1.f3[2], i + 11) + y1.f3[3] * partials(o1.f3[3], i + 11) + y1.f4[1] * partials(o1.f4[1], i + 11) + y1.f4[2] * partials(o1.f4[2], i + 11) + y1.f4[3] * partials(o1.f4[3], i + 11),
        ),
        false,
    )
end

function combine_dual_Atom(y1::SVector{3, T}, o1::SVector{3, Dual{Nothing, T, P}}, i::Integer, j::Integer, k::Integer, l::Integer) where {T, P}
    ps1, ps2, ps3 = partials(o1[1]), partials(o1[2]), partials(o1[3])
    Atom(
        0,
        y1[1] * ps1[i] + y1[2] * ps2[i] + y1[3] * ps3[i],
        y1[1] * ps1[j] + y1[2] * ps2[j] + y1[3] * ps3[j],
        y1[1] * ps1[k] + y1[2] * ps2[k] + y1[3] * ps3[k],
        y1[1] * ps1[l] + y1[2] * ps2[l] + y1[3] * ps3[l],
        false,
    )
end

# For force_nounit
@inline function Zygote.broadcast_forward(f,
                                            arg1,
                                            arg2::AbstractArray{SVector{D, T}},
                                            arg3::AbstractArray{SVector{D, T}},
                                            arg4::AbstractArray{<:Atom},
                                            arg5::AbstractArray{<:Atom},
                                            arg6::Tuple{SVector{D, T}},
                                            arg7::Base.RefValue{<:Unitful.FreeUnits},
                                            arg8) where {D, T}
    out = dual_function_force_broadcast(f).(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8)
    y = map(x -> value.(x), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg2)
        darg1 = unbroadcast(arg1, broadcast(combine_dual_PairwiseInteraction, ȳ, out, 1))
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(sum_partials(o1, y1,  7),
                                sum_partials(o1, y1,  8), sum_partials(o1, y1,  9)), ȳ, out))
        darg3 = unbroadcast(arg3, broadcast((y1, o1) -> SVector{D, T}(sum_partials(o1, y1, 10),
                                sum_partials(o1, y1, 11), sum_partials(o1, y1, 12)), ȳ, out))
        darg4 = unbroadcast(arg4, broadcast(combine_dual_Atom, ȳ, out, 13, 14, 15, 16))
        darg5 = unbroadcast(arg5, broadcast(combine_dual_Atom, ȳ, out, 17, 18, 19, 20))
        darg6 = unbroadcast(arg6, broadcast((y1, o1) -> SVector{D, T}(sum_partials(o1, y1, 21),
                                sum_partials(o1, y1, 22), sum_partials(o1, y1, 23)), ȳ, out))
        darg7 = nothing
        darg8 = nothing
        return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6, darg7, darg8)
    end
    return y, bc_fwd_back
end

# For force
@inline function Zygote.broadcast_forward(f,
                                            arg1::AbstractArray{<:SpecificInteraction},
                                            arg2::AbstractArray{SVector{D, T}},
                                            arg3::AbstractArray{SVector{D, T}},
                                            arg4::Tuple{SVector{D, T}}) where {D, T}
    out = dual_function_specific_2_atoms(f).(arg1, arg2, arg3, arg4)
    y = broadcast(o1 -> SpecificForce2Atoms{D, T}(value.(o1.f1), value.(o1.f2)), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        darg1 = unbroadcast(arg1, broadcast(combine_dual_SpecificInteraction, arg1, ȳ, out, 1))
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 3) + sum_partials(o1.f2, y1.f2, 3),
                    sum_partials(o1.f1, y1.f1, 4) + sum_partials(o1.f2, y1.f2, 4),
                    sum_partials(o1.f1, y1.f1, 5) + sum_partials(o1.f2, y1.f2, 5)),
                    ȳ, out))
        darg3 = unbroadcast(arg3, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 6) + sum_partials(o1.f2, y1.f2, 6),
                    sum_partials(o1.f1, y1.f1, 7) + sum_partials(o1.f2, y1.f2, 7),
                    sum_partials(o1.f1, y1.f1, 8) + sum_partials(o1.f2, y1.f2, 8)),
                    ȳ, out))
        darg4 = unbroadcast(arg4, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1,  9) + sum_partials(o1.f2, y1.f2,  9),
                    sum_partials(o1.f1, y1.f1, 10) + sum_partials(o1.f2, y1.f2, 10),
                    sum_partials(o1.f1, y1.f1, 11) + sum_partials(o1.f2, y1.f2, 11)),
                    ȳ, out))
        return (nothing, nothing, darg1, darg2, darg3, darg4)
    end
    return y, bc_fwd_back
end

# For force
@inline function Zygote.broadcast_forward(f,
                                            arg1::AbstractArray{<:SpecificInteraction},
                                            arg2::AbstractArray{SVector{D, T}},
                                            arg3::AbstractArray{SVector{D, T}},
                                            arg4::AbstractArray{SVector{D, T}},
                                            arg5::Tuple{SVector{D, T}}) where {D, T}
    out = dual_function_specific_3_atoms(f).(arg1, arg2, arg3, arg4, arg5)
    y = broadcast(o1 -> SpecificForce3Atoms{D, T}(value.(o1.f1), value.(o1.f2), value.(o1.f3)), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        darg1 = unbroadcast(arg1, broadcast(combine_dual_SpecificInteraction, arg1, ȳ, out, 1))
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 3) + sum_partials(o1.f2, y1.f2, 3) + sum_partials(o1.f3, y1.f3, 3),
                    sum_partials(o1.f1, y1.f1, 4) + sum_partials(o1.f2, y1.f2, 4) + sum_partials(o1.f3, y1.f3, 4),
                    sum_partials(o1.f1, y1.f1, 5) + sum_partials(o1.f2, y1.f2, 5) + sum_partials(o1.f3, y1.f3, 5)),
                    ȳ, out))
        darg3 = unbroadcast(arg3, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 6) + sum_partials(o1.f2, y1.f2, 6) + sum_partials(o1.f3, y1.f3, 6),
                    sum_partials(o1.f1, y1.f1, 7) + sum_partials(o1.f2, y1.f2, 7) + sum_partials(o1.f3, y1.f3, 7),
                    sum_partials(o1.f1, y1.f1, 8) + sum_partials(o1.f2, y1.f2, 8) + sum_partials(o1.f3, y1.f3, 8)),
                    ȳ, out))
        darg4 = unbroadcast(arg4, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1,  9) + sum_partials(o1.f2, y1.f2,  9) + sum_partials(o1.f3, y1.f3,  9),
                    sum_partials(o1.f1, y1.f1, 10) + sum_partials(o1.f2, y1.f2, 10) + sum_partials(o1.f3, y1.f3, 10),
                    sum_partials(o1.f1, y1.f1, 11) + sum_partials(o1.f2, y1.f2, 11) + sum_partials(o1.f3, y1.f3, 11)),
                    ȳ, out))
        darg5 = unbroadcast(arg5, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 12) + sum_partials(o1.f2, y1.f2, 12) + sum_partials(o1.f3, y1.f3, 12),
                    sum_partials(o1.f1, y1.f1, 13) + sum_partials(o1.f2, y1.f2, 13) + sum_partials(o1.f3, y1.f3, 13),
                    sum_partials(o1.f1, y1.f1, 14) + sum_partials(o1.f2, y1.f2, 14) + sum_partials(o1.f3, y1.f3, 14)),
                    ȳ, out))
        return (nothing, nothing, darg1, darg2, darg3, darg4, darg5)
    end
    return y, bc_fwd_back
end

# For force
@inline function Zygote.broadcast_forward(f,
                                            arg1::AbstractArray{<:SpecificInteraction},
                                            arg2::AbstractArray{SVector{D, T}},
                                            arg3::AbstractArray{SVector{D, T}},
                                            arg4::AbstractArray{SVector{D, T}},
                                            arg5::AbstractArray{SVector{D, T}},
                                            arg6::Tuple{SVector{D, T}}) where {D, T}
    out = dual_function_specific_4_atoms(f).(arg1, arg2, arg3, arg4, arg5, arg6)
    y = broadcast(o1 -> SpecificForce4Atoms{D, T}(value.(o1.f1), value.(o1.f2), value.(o1.f3), value.(o1.f4)), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        darg1 = unbroadcast(arg1, broadcast(combine_dual_SpecificInteraction, arg1, ȳ, out, 1))
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 13) + sum_partials(o1.f2, y1.f2, 13) + sum_partials(o1.f3, y1.f3, 13) + sum_partials(o1.f4, y1.f4, 13),
                    sum_partials(o1.f1, y1.f1, 14) + sum_partials(o1.f2, y1.f2, 14) + sum_partials(o1.f3, y1.f3, 14) + sum_partials(o1.f4, y1.f4, 14),
                    sum_partials(o1.f1, y1.f1, 15) + sum_partials(o1.f2, y1.f2, 15) + sum_partials(o1.f3, y1.f3, 15) + sum_partials(o1.f4, y1.f4, 15)),
                    ȳ, out))
        darg3 = unbroadcast(arg3, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 16) + sum_partials(o1.f2, y1.f2, 16) + sum_partials(o1.f3, y1.f3, 16) + sum_partials(o1.f4, y1.f4, 16),
                    sum_partials(o1.f1, y1.f1, 17) + sum_partials(o1.f2, y1.f2, 17) + sum_partials(o1.f3, y1.f3, 17) + sum_partials(o1.f4, y1.f4, 17),
                    sum_partials(o1.f1, y1.f1, 18) + sum_partials(o1.f2, y1.f2, 18) + sum_partials(o1.f3, y1.f3, 18) + sum_partials(o1.f4, y1.f4, 18)),
                    ȳ, out))
        darg4 = unbroadcast(arg4, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 19) + sum_partials(o1.f2, y1.f2, 19) + sum_partials(o1.f3, y1.f3, 19) + sum_partials(o1.f4, y1.f4, 19),
                    sum_partials(o1.f1, y1.f1, 20) + sum_partials(o1.f2, y1.f2, 20) + sum_partials(o1.f3, y1.f3, 20) + sum_partials(o1.f4, y1.f4, 20),
                    sum_partials(o1.f1, y1.f1, 21) + sum_partials(o1.f2, y1.f2, 21) + sum_partials(o1.f3, y1.f3, 21) + sum_partials(o1.f4, y1.f4, 21)),
                    ȳ, out))
        darg5 = unbroadcast(arg5, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 22) + sum_partials(o1.f2, y1.f2, 22) + sum_partials(o1.f3, y1.f3, 22) + sum_partials(o1.f4, y1.f4, 22),
                    sum_partials(o1.f1, y1.f1, 23) + sum_partials(o1.f2, y1.f2, 23) + sum_partials(o1.f3, y1.f3, 23) + sum_partials(o1.f4, y1.f4, 23),
                    sum_partials(o1.f1, y1.f1, 24) + sum_partials(o1.f2, y1.f2, 24) + sum_partials(o1.f3, y1.f3, 24) + sum_partials(o1.f4, y1.f4, 24)),
                    ȳ, out))
        darg6 = unbroadcast(arg6, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.f1, y1.f1, 25) + sum_partials(o1.f2, y1.f2, 25) + sum_partials(o1.f3, y1.f3, 25) + sum_partials(o1.f4, y1.f4, 25),
                    sum_partials(o1.f1, y1.f1, 26) + sum_partials(o1.f2, y1.f2, 26) + sum_partials(o1.f3, y1.f3, 26) + sum_partials(o1.f4, y1.f4, 26),
                    sum_partials(o1.f1, y1.f1, 27) + sum_partials(o1.f2, y1.f2, 27) + sum_partials(o1.f3, y1.f3, 27) + sum_partials(o1.f4, y1.f4, 27)),
                    ȳ, out))
        return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6)
    end
    return y, bc_fwd_back
end

# For born_radii_loop
@inline function Zygote.broadcast_forward(f,
                                            arg1::AbstractArray{SVector{D, T}},
                                            arg2::AbstractArray{SVector{D, T}},
                                            arg3::AbstractArray{T},
                                            arg4::AbstractArray{T},
                                            arg5::Tuple{T},
                                            arg6::Tuple{SVector{D, T}}) where {D, T}
    out = dual_function_born_radii_loop(f).(arg1, arg2, arg3, arg4, arg5, arg6)
    y = value.(out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        darg1 = unbroadcast(arg1, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 1) * y1,
                    partials(o1, 2) * y1, partials(o1, 3) * y1), ȳ, out))
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 4) * y1,
                    partials(o1, 5) * y1, partials(o1, 6) * y1), ȳ, out))
        darg3 = unbroadcast(arg3, broadcast((y1, o1) -> partials(o1, 7) * y1, ȳ, out))
        darg4 = unbroadcast(arg4, broadcast((y1, o1) -> partials(o1, 8) * y1, ȳ, out))
        darg5 = nothing
        darg6 = unbroadcast(arg6, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 9) * y1,
                    partials(o1, 10) * y1, partials(o1, 11) * y1), ȳ, out))
        return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6)
    end
    return y, bc_fwd_back
end

# For gb_force_loop_1
@inline function Zygote.broadcast_forward(f,
                                            arg1::AbstractArray{SVector{D, T}},
                                            arg2::AbstractArray{SVector{D, T}},
                                            arg3::AbstractArray{Int},
                                            arg4::AbstractArray{Int},
                                            arg5::AbstractArray{T},
                                            arg6::AbstractArray{T},
                                            arg7::AbstractArray{T},
                                            arg8::AbstractArray{T},
                                            arg9::Tuple{T},
                                            arg10::Tuple{T},
                                            arg11::Tuple{SVector{D, T}}) where {D, T}
    out = dual_function_gb_force_loop_1(f).(arg1, arg2, arg3, arg4, arg5, arg6, arg7,
                                            arg8, arg9, arg10, arg11)
    y = broadcast(o1 -> ForceLoopResult1{T, SVector{D, T}}(value(o1.bi), value(o1.bj),
                                            value.(o1.fi), value.(o1.fj)), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        darg1 = unbroadcast(arg1, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.fi, y1.fi, 1) + sum_partials(o1.fj, y1.fj, 1) + partials(o1.bi, 1) * y1.bi + partials(o1.bj, 1) * y1.bj,
                    sum_partials(o1.fi, y1.fi, 2) + sum_partials(o1.fj, y1.fj, 2) + partials(o1.bi, 2) * y1.bi + partials(o1.bj, 2) * y1.bj,
                    sum_partials(o1.fi, y1.fi, 3) + sum_partials(o1.fj, y1.fj, 3) + partials(o1.bi, 3) * y1.bi + partials(o1.bj, 3) * y1.bj),
                    ȳ, out))
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.fi, y1.fi, 4) + sum_partials(o1.fj, y1.fj, 4) + partials(o1.bi, 4) * y1.bi + partials(o1.bj, 4) * y1.bj,
                    sum_partials(o1.fi, y1.fi, 5) + sum_partials(o1.fj, y1.fj, 5) + partials(o1.bi, 5) * y1.bi + partials(o1.bj, 5) * y1.bj,
                    sum_partials(o1.fi, y1.fi, 6) + sum_partials(o1.fj, y1.fj, 6) + partials(o1.bi, 6) * y1.bi + partials(o1.bj, 6) * y1.bj),
                    ȳ, out))
        darg3 = nothing
        darg4 = nothing
        darg5 = unbroadcast(arg5, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi,  7) + sum_partials(o1.fj, y1.fj,  7) + partials(o1.bi,  7) * y1.bi + partials(o1.bj,  7) * y1.bj, ȳ, out))
        darg6 = unbroadcast(arg6, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi,  8) + sum_partials(o1.fj, y1.fj,  8) + partials(o1.bi,  8) * y1.bi + partials(o1.bj,  8) * y1.bj, ȳ, out))
        darg7 = unbroadcast(arg7, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi,  9) + sum_partials(o1.fj, y1.fj,  9) + partials(o1.bi,  9) * y1.bi + partials(o1.bj,  9) * y1.bj, ȳ, out))
        darg8 = unbroadcast(arg8, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 10) + sum_partials(o1.fj, y1.fj, 10) + partials(o1.bi, 10) * y1.bi + partials(o1.bj, 10) * y1.bj, ȳ, out))
        darg9 = nothing
        darg10 = unbroadcast(arg10, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 11) + sum_partials(o1.fj, y1.fj, 11) + partials(o1.bi, 11) * y1.bi + partials(o1.bj, 11) * y1.bj, ȳ, out))
        darg11 = unbroadcast(arg11, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.fi, y1.fi, 12) + sum_partials(o1.fj, y1.fj, 12) + partials(o1.bi, 12) * y1.bi + partials(o1.bj, 12) * y1.bj,
                    sum_partials(o1.fi, y1.fi, 13) + sum_partials(o1.fj, y1.fj, 13) + partials(o1.bi, 13) * y1.bi + partials(o1.bj, 13) * y1.bj,
                    sum_partials(o1.fi, y1.fi, 14) + sum_partials(o1.fj, y1.fj, 14) + partials(o1.bi, 14) * y1.bi + partials(o1.bj, 14) * y1.bj),
                    ȳ, out))
        return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6, darg7,
                darg8, darg9, darg10, darg11)
    end
    return y, bc_fwd_back
end

# For gb_force_loop_2
@inline function Zygote.broadcast_forward(f,
                                            arg1::AbstractArray{SVector{D, T}},
                                            arg2::AbstractArray{SVector{D, T}},
                                            arg3::AbstractArray{T},
                                            arg4::AbstractArray{T},
                                            arg5::AbstractArray{T},
                                            arg6::Tuple{T},
                                            arg7::Tuple{SVector{D, T}}) where {D, T}
    out = dual_function_gb_force_loop_2(f).(arg1, arg2, arg3, arg4, arg5, arg6, arg7)
    y = broadcast(o1 -> ForceLoopResult2{SVector{D, T}}(value.(o1.fi), value.(o1.fj)), out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        darg1 = unbroadcast(arg1, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.fi, y1.fi, 1) + sum_partials(o1.fj, y1.fj, 1),
                    sum_partials(o1.fi, y1.fi, 2) + sum_partials(o1.fj, y1.fj, 2),
                    sum_partials(o1.fi, y1.fi, 3) + sum_partials(o1.fj, y1.fj, 3)),
                    ȳ, out))
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.fi, y1.fi, 4) + sum_partials(o1.fj, y1.fj, 4),
                    sum_partials(o1.fi, y1.fi, 5) + sum_partials(o1.fj, y1.fj, 5),
                    sum_partials(o1.fi, y1.fi, 6) + sum_partials(o1.fj, y1.fj, 6)),
                    ȳ, out))
        darg3 = unbroadcast(arg3, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 7) + sum_partials(o1.fj, y1.fj, 7), ȳ, out))
        darg4 = unbroadcast(arg4, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 8) + sum_partials(o1.fj, y1.fj, 8), ȳ, out))
        darg5 = unbroadcast(arg5, broadcast((y1, o1) -> sum_partials(o1.fi, y1.fi, 9) + sum_partials(o1.fj, y1.fj, 9), ȳ, out))
        darg6 = nothing
        darg7 = unbroadcast(arg7, broadcast((y1, o1) -> SVector{D, T}(
                    sum_partials(o1.fi, y1.fi, 10) + sum_partials(o1.fj, y1.fj, 10),
                    sum_partials(o1.fi, y1.fi, 11) + sum_partials(o1.fj, y1.fj, 11),
                    sum_partials(o1.fi, y1.fi, 12) + sum_partials(o1.fj, y1.fj, 12)),
                    ȳ, out))
        return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6, darg7)
    end
    return y, bc_fwd_back
end

# For gb_energy_loop
@inline function Zygote.broadcast_forward(f,
                                            arg1::AbstractArray{SVector{D, T}},
                                            arg2::AbstractArray{SVector{D, T}},
                                            arg3::AbstractArray{Int},
                                            arg4::AbstractArray{Int},
                                            arg5::AbstractArray{T},
                                            arg6::AbstractArray{T},
                                            arg7::AbstractArray{T},
                                            arg8::AbstractArray{T},
                                            arg9::AbstractArray{T},
                                            arg10::Tuple{T},
                                            arg11::Tuple{T},
                                            arg12::Tuple{T},
                                            arg13::Tuple{T},
                                            arg14::Tuple{T},
                                            arg15::Tuple{Bool},
                                            arg16::Tuple{SVector{D, T}}) where {D, T}
    out = dual_function_gb_energy_loop(f).(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9,
                                            arg10, arg11, arg12, arg13, arg14, arg15, arg16)
    y = value.(out)
    function bc_fwd_back(ȳ_in)
        ȳ = modify_grad(ȳ_in, arg1)
        darg1 = unbroadcast(arg1, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 1) * y1,
                    partials(o1, 2) * y1, partials(o1, 3) * y1), ȳ, out))
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 4) * y1,
                    partials(o1, 5) * y1, partials(o1, 6) * y1), ȳ, out))
        darg3 = nothing
        darg4 = nothing
        darg5 = unbroadcast(arg5, broadcast((y1, o1) -> partials(o1,  7) * y1, ȳ, out))
        darg6 = unbroadcast(arg6, broadcast((y1, o1) -> partials(o1,  8) * y1, ȳ, out))
        darg7 = unbroadcast(arg7, broadcast((y1, o1) -> partials(o1,  9) * y1, ȳ, out))
        darg8 = unbroadcast(arg8, broadcast((y1, o1) -> partials(o1, 10) * y1, ȳ, out))
        darg9 = unbroadcast(arg9, broadcast((y1, o1) -> partials(o1, 11) * y1, ȳ, out))
        darg10 = nothing
        darg11 = unbroadcast(arg11, broadcast((y1, o1) -> partials(o1, 12) * y1, ȳ, out))
        darg12 = unbroadcast(arg12, broadcast((y1, o1) -> partials(o1, 13) * y1, ȳ, out))
        darg13 = unbroadcast(arg13, broadcast((y1, o1) -> partials(o1, 14) * y1, ȳ, out))
        darg14 = unbroadcast(arg14, broadcast((y1, o1) -> partials(o1, 15) * y1, ȳ, out))
        darg15 = nothing
        darg16 = unbroadcast(arg16, broadcast((y1, o1) -> SVector{D, T}(partials(o1, 16) * y1,
                    partials(o1, 17) * y1, partials(o1, 18) * y1), ȳ, out))
        return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6, darg7, darg8, darg9,
                darg10, darg11, darg12, darg13, darg14, darg15, darg16)
    end
    return y, bc_fwd_back
end

@inline function Zygote.broadcast_forward(f::typeof(getf1),
                                            arg1::AbstractArray{<:SpecificForce2Atoms}) where {D, T}
    return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (f1=y1, f2=zero(y1)), ȳ)))
end

@inline function Zygote.broadcast_forward(f::typeof(getf1),
                                            arg1::AbstractArray{<:SpecificForce3Atoms}) where {D, T}
    return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (f1=y1, f2=zero(y1), f3=zero(y1)), ȳ)))
end

@inline function Zygote.broadcast_forward(f::typeof(getf1),
                                            arg1::AbstractArray{<:SpecificForce4Atoms}) where {D, T}
    return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (f1=y1, f2=zero(y1), f3=zero(y1), f4=zero(y1)), ȳ)))
end

@inline function Zygote.broadcast_forward(f::typeof(getf2),
                                            arg1::AbstractArray{<:SpecificForce2Atoms}) where {D, T}
    return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (f1=zero(y1), f2=y1), ȳ)))
end

@inline function Zygote.broadcast_forward(f::typeof(getf2),
                                            arg1::AbstractArray{<:SpecificForce3Atoms}) where {D, T}
    return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (f1=zero(y1), f2=y1, f3=zero(y1)), ȳ)))
end

@inline function Zygote.broadcast_forward(f::typeof(getf2),
                                            arg1::AbstractArray{<:SpecificForce4Atoms}) where {D, T}
    return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (f1=zero(y1), f2=y1, f3=zero(y1), f4=zero(y1)), ȳ)))
end

@inline function Zygote.broadcast_forward(f::typeof(getf3),
                                            arg1::AbstractArray{<:SpecificForce3Atoms}) where {D, T}
    return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (f1=zero(y1), f2=zero(y1), f3=y1), ȳ)))
end

@inline function Zygote.broadcast_forward(f::typeof(getf3),
                                            arg1::AbstractArray{<:SpecificForce4Atoms}) where {D, T}
    return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (f1=zero(y1), f2=zero(y1), f3=y1, f4=zero(y1)), ȳ)))
end

@inline function Zygote.broadcast_forward(f::typeof(getf4),
                                            arg1::AbstractArray{<:SpecificForce4Atoms}) where {D, T}
    return f.(arg1), ȳ -> (nothing, nothing, unbroadcast(arg1, broadcast(y1 -> (f1=zero(y1), f2=zero(y1), f3=zero(y1), f4=y1), ȳ)))
end

# Use fast broadcast path on CPU
for op in (:+, :-, :*, :/, :mass, :charge, :remove_molar, :ustrip, :ustrip_vec, :wrap_coords_vec,
            :getf1, :getf2, :getf3, :getf4, :born_radii_loop, :gb_force_loop_1, :gb_force_loop_2,
            :gb_energy_loop)
    @eval Zygote.@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::typeof($op), args...) = Zygote.broadcast_forward(f, args...)
    # Avoid ambiguous dispatch
    @eval Zygote.@adjoint Broadcast.broadcasted(::CUDA.AbstractGPUArrayStyle  , f::typeof($op), args...) = Zygote.broadcast_forward(f, args...)
end

# Interactions not specified here run on the slow path on CPU
@eval Zygote.@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::typeof(force_nounit), inters::Tuple{Tuple{X, Y}}, args...) where {X <: LennardJones, Y <: CoulombReactionField} = Zygote.broadcast_forward(f, inters, args...)

@eval Zygote.@adjoint Broadcast.broadcasted(::CUDA.AbstractGPUArrayStyle  , f::typeof(force_nounit), inters::Tuple{Tuple{X, Y}}, args...) where {X <: LennardJones, Y <: CoulombReactionField} = Zygote.broadcast_forward(f, inters, args...)

@eval Zygote.@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::typeof(force), inters::Vector{HarmonicBond{D, K}}      , args...) where {D, K}    = Zygote.broadcast_forward(f, inters, args...)
@eval Zygote.@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::typeof(force), inters::Vector{HarmonicAngle{D, K}}     , args...) where {D, K}    = Zygote.broadcast_forward(f, inters, args...)
@eval Zygote.@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::typeof(force), inters::Vector{PeriodicTorsion{N, T, E}}, args...) where {N, T, E} = Zygote.broadcast_forward(f, inters, args...)

@eval Zygote.@adjoint Broadcast.broadcasted(::CUDA.AbstractGPUArrayStyle  , f::typeof(force), inters::Vector{HarmonicBond{D, K}}      , args...) where {D, K}    = Zygote.broadcast_forward(f, inters, args...)
@eval Zygote.@adjoint Broadcast.broadcasted(::CUDA.AbstractGPUArrayStyle  , f::typeof(force), inters::Vector{HarmonicAngle{D, K}}     , args...) where {D, K}    = Zygote.broadcast_forward(f, inters, args...)
@eval Zygote.@adjoint Broadcast.broadcasted(::CUDA.AbstractGPUArrayStyle  , f::typeof(force), inters::Vector{PeriodicTorsion{N, T, E}}, args...) where {N, T, E} = Zygote.broadcast_forward(f, inters, args...)


text:
module DecisionMakingUtils

using Base: Float64
using LinearAlgebra
using OnlineStats
import OnlineStats: fit!
import Base: length, eltype

export FourierBasis, FourierBasisBuffer, ConcatentateBasis
export LinearNormalization, ZeroOneNormalization, PosNegNormalization, GaussianNormalization
export gaussian_stats, extrema_stats, fit!
export BufferedFunction
export update!

include("basis/basis.jl")
include("bufferedfun.jl")
# include("normalization.jl")
include("normalization_constant.jl")
include("stats.jl")

end


text:
using Pythia
using Test

### Mean Forecast Tests ###
y = [1.0, 2.0, 3.0, 4.0, 5.0]
mdl = MeanForecast(y, h = 5, level = [80, 95])
fittedMdl = fit(mdl)
results = predict(fittedMdl)

@test results.fittedvalues == [1.0, 2.0, 3.0, 4.0, 5.0, 3.0, 3.0, 3.0, 3.0, 3.0]
### Will add tests for prediction intervals in later commits.

### Naive Forecast Tests ###
y = [1.0, 2.0, 3.0, 4.0, 5.0]
mdl = NaiveForecast(y, h = 5, level = [80, 95])
fittedMdl = fit(mdl)
results = predict(fittedMdl)

@test results.fittedvalues == [1.0, 2.0, 3.0, 4.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0]
### Will add tests for prediction intervals in later commits.

### Seasonal Naive Model ###