module NonMeansTestedBenefits
#
# This module provides routines for benefits which are (mostly) not means-tested benefits,
# both current and historic. See [HistoricBenefits.jl] for the imputation routines used here.
#
    using Base: has_tight_type
    using Dates
    using Dates: Date, now, TimeType, Year
    using ScottishTaxBenefitModel

    using .Definitions
    using .STBIncomes

    using .Utils: nearest

    using .ModelHousehold:
        BenefitUnit,
        Household,
        Person,
        get_benefit_units,
        has_income,
        num_children

    using .BenefitGenerosity: change_status

    using .Intermediate:
        has_limited_capactity_for_work_activity,
        has_limited_capactity_for_work,
        make_recipient

    using .STBParameters:
        AgeLimits,
        AttendanceAllowance,
        BereavementSupport,
        CarersAllowance,
        ChildBenefit,
        ContributoryESA,
        DisabilityLivingAllowance,
        HoursLimits,
        JobSeekersAllowance,
        MaternityAllowance,
        NonMeansTestedSys,
        PersonalIndependencePayment,
        RetirementPension,
        WidowsPensions,
        reached_state_pension_age

    using .Results:
        BenefitUnitResult,
        HouseholdResult,
        IndividualResult,
        LMTIncomes

    export
        calc_child_benefit!,
        calc_dla,
        calc_pip,
        calc_post_tax_non_means_tested!,
        calc_pre_tax_non_means_tested!,
        calc_state_pension,
        calc_widows_benefits

    """
    Child Benefit - this has to be done *after* income tax, so we have
    income tax `total_income` for each adult in the `BenefitUnitResult` struct.
    """
    function calc_child_benefit!(
        bures :: BenefitUnitResult,
        bu    :: BenefitUnit,
        cb    :: ChildBenefit{T} ) :: T where T
        c = zero(T)
        #= FIXME something to check the exact type of children -
        for now, assume the BU allocation has got this right.
        for pid in bu.children
            #
        end
        =#
        if cb.abolished
            return c
        end
        nc = num_children(bu)
        if nc == 0
            return 0.0
        end
        c += cb.first_child
        if nc > 1
            c += (nc-1)*cb.other_children
        end
        recipient = make_recipient( bu, CHILD_BENEFIT )
        # guardian's allowance
        bures.pers[recipient].income[GUARDIANS_ALLOWANCE] = 0.0

        # println( bu.people[bu.head].relationships)

        if c > 0 # fixme not quite right - qualify for CB for each child
            for cp in bu.children
                # this checks if anyone in the BU has a parent-like
                # relationship to the child
                is_guardian = true
                for pid in bu.adults
                    if ! (bu.people[pid].relationships[cp] in [Grand_parent, Other_relative, Other_non_relative])
                        is_guardian = false
                        break;
                    end
                end
                if is_guardian
                    # FIXME we could also guess approx number of kids
                    # qualifying for GA from receipt in HistoricBenefits || has_income( bu, guardians_allowance )
                    bures.pers[recipient].income[GUARDIANS_ALLOWANCE] +=
                        cb.guardians_allowance
                end
            end
        end # guardian's allowance
        # cb but not guardians withdrawn with high incomes
        # high income thing - highest of the BU's *individual* income; see cpag ch.27
        max_inc = zero(T)
        for pid in bu.adults
            max_inc = max( max_inc, bures.pers[pid].it.total_income )
            if bu.people[pid].sex == Female
                recipient = pid
            end
        end
        if max_inc > cb.high_income_thresh
            # this should really be done in steps £1 for every £100, but since
            # everything is weekly here we'll just multiply
            withdrawn = cb.withdrawal*(max_inc - cb.high_income_thresh)
            c = max(0.0, c-withdrawn)
        end
        bures.pers[recipient].income[CHILD_BENEFIT] = c
    end

    function calc_widows_benefits(
        pers :: Person{T},
        has_kids :: Bool,
        bp :: BereavementSupport{T},
        wp :: WidowsPensions{T}) :: T where T
        wid = zero(T)
        if wp.abolished || bp.abolished
            return wid
        end
        #
        # We don't know when someone was widowed
        # so we rely on this. FIXME Obviously this gets progressively worse as time goes on and
        # we keep using the old years of FRS.
        #
        # new-style: payable for 18 months at a flat rate
        # CHECK 2829 2017 why is that bereavment_support?
        if pers.bereavement_type in [bereavement_allowance,bereavement_support]
            # payable for 18 months so we'll allocate
            # 2/3rds of the weekly value of the lump-sum
            if has_kids
                wid = bp.higher + bp.lump_sum_higher*2/3
            else
                wid = bp.lower + bp.lump_sum_lower*2/3
            end
        elseif pers.bereavement_type == widowed_parents # At least 3 years ago so no
            # need to worry about lump-sums; just scale
            # the standard rate by the ratio of
            # their receipt to the standard rate at the time
            # of interview. CHECK hhld 477 2015 for an example
            # of someone with the lump sum but no standard benefit
            # so no ratio; for us, just assign the standard
            # rate in those cases.
            #
            # println( "pers.benefit_ratios $(pers.benefit_ratios)")
            if haskey( pers.benefit_ratios, bereavement_allowance_or_widowed_parents_allowance_or_bereavement )
                wid = pers.benefit_ratios[bereavement_allowance_or_widowed_parents_allowance_or_bereavement]*
                    wp.standard_rate
            else
                wid = wp.standard_rate
            end
        else
            # check we're not missing anyone
            @assert ! haskey( pers.benefit_ratios, bereavement_allowance_or_widowed_parents_allowance_or_bereavement)
        end
        return wid
    end

    """
    PIP and DLA (below) rely on all the types being sorted out earlier
    either in some kind of probit or by being inferred from receipts (see [HistoricBenefits.jl] for
    what we have so far).
    """
    function calc_pip(
        pers :: Person{T},
        pip  :: PersonalIndependencePayment{T}) :: Tuple{T,T} where T
        pl = zero(T)
        pm = zero(T)
        daily_type = pers.pip_daily_living_type
        mob_type = pers.pip_mobility_type
        if pip.abolished
            return (pl, pm )
        end
        daily_type = change_status(
            candidates=pip.dl_candidates,
            pid=pers.pid,
            change=pip.extra_people,
            choices=[standard_pip,enhanced_pip],
            current_value=daily_type,
            disqual_value=no_pip )
        mob_type = change_status(
            candidates=pip.mobility_candidates,
            pid=pers.pid,
            choices=[standard_pip,enhanced_pip],
            change=pip.extra_people,
            current_value=mob_type,
            disqual_value=no_pip )

        # fixme consistent names dl->daily living etc.
        if daily_type == standard_pip
            pl = pip.dl_standard
        elseif daily_type == enhanced_pip
            pl = pip.dl_enhanced
        end
        if mob_type == standard_pip
            pm = pip.mobility_standard
        elseif mob_type == enhanced_pip
            pm = pip.mobility_enhanced
        end
        return (pl, pm )
    end # pip calc

    function calc_attendance_allowance(
        pers :: Person{T},
        aa  :: AttendanceAllowance{T},
         ) :: T where T
        a =zero(T)
        if aa.abolished
            return a
        end
        at = pers.attendance_allowance_type
        at = change_status(
            candidates=aa.candidates,
            pid=pers.pid,
            choices=[high,low],
            change=aa.extra_people,
            current_value=at,
            disqual_value=missing_lmh  )
        if at == missing_lmh
            a = zero(T)
        elseif at == high
            a = aa.higher
        else
            a = aa.lower;
        end
        return a
    end

    function calc_dla(
        pers :: Person{T},
        dla  :: DisabilityLivingAllowance{T} ) :: Tuple{T,T} where T
        dc = zero(T)
        dm = zero(T)
        if dla.abolished
            return (dc,dm)
        end
        dla_s = pers.dla_self_care_type
        dla_m = pers.dla_mobility_type
        # FIXME we use the same list for both mob and self
        # I think because of small sample size (kids only)
        dla_s = change_status(
            candidates=dla.candidates,
            pid=pers.pid,
            choices=[high,mid,low],
            change=dla.extra_people,
            current_value=dla_s,
            disqual_value=missing_lmh  )
        dla_m = change_status(
            candidates=dla.candidates,
            pid=pers.pid,
            choices=[high,mid,low],
            change=dla.extra_people,
            current_value=dla_m,
            disqual_value=missing_lmh  )
        # FIXME make all these names constisent (mid/middle,care->self_care etc.)
        if dla_s == high
            dc = dla.care_high
        elseif dla_s == mid
            dc = dla.care_middle
        elseif dla_s == low
            dc = dla.care_low
        end
        if dla_m == high
            dm = dla.mob_high
        elseif dla_m in (low,mid)
            dm = dla.mob_low
        end
        # println( "setting DLA as $dc, $dm")
        return (dc,dm);
    end # dla calc

    function calc_state_pension(
        pers :: Person{T},
        rp :: RetirementPension{T},
        age_limits :: AgeLimits ) :: T where T
        pen = zero(T)
        if rp.abolished
            return pen
        end
        if reached_state_pension_age(
            age_limits,
            pers.age,
            pers.sex )
            # so, was someone over pension age *before* the new state pension?
            # if so, use the proportion of the `class_a` they seemed to be on at
            # the time, times the current `class_a`.
            if reached_state_pension_age(
                # old style
                age_limits,
                pers.age,
                pers.sex,
                age_limits.savings_credit_to_new_state_pension )
                # println( "pers.benefit_ratios $(pers.benefit_ratios)")
                if ! haskey( pers.benefit_ratios, state_pension )
                    ratio =  pers.age < 70 ? 0.0 : 1.0 # kinda crude deferrment
                else
                    ratio = pers.benefit_ratios[state_pension]
                end
                pen = ratio*rp.cat_a
            else
                #
                # new style
                pen = rp.new_state_pension
            end
        end
        return pen
    end

    """
    FIXME Model this properly
    """
    function calc_esa(
        pers :: Person{T},
        esa  :: ContributoryESA{T}) :: T where T
        e = zero(T)
        if esa.abolished
            return e
        end
        if pers.esa_type == contributory_jsa
            if pers.age < 25
                # FIXME not quite right since
                # could be past assessment stage;
                # maybe check time out of work?
                e = esa.assessment_u25
            else
                e = esa.main
            end
            if has_limited_capactity_for_work_activity( pers )
                e += esa.support
            end
        end
        return e
    end

    function calc_maternity_allowance(
        pers :: Person{T},
        ma :: MaternityAllowance ) :: T where T
        m = zero(T)
        if ma.abolished
            return m
        end
        # fixme the design means you should never have to check the incomes dict here
        if has_income( pers, maternity_allowance )
            m = ma.rate
        end
        return m
    end

    function calc_carers_allowance(
        pers :: Person{T},
        pres :: IndividualResult{T},
        carers :: CarersAllowance{T}) :: T where T
        c = zero(T)
        if carers.abolished
            return c
        end
        earnings :: T = isum(
            pres.income,
            carers.earnings;
            deducted=carers.deductions )
        # println( "earnings=$earnings carers.gainful_employment_min=$(carers.gainful_employment_min)")
        if pers.hours_of_care_given >= carers.hours &&
            earnings < carers.gainful_employment_min
            c = carers.allowance
        end
        return c
    end

    function calc_jsa(
        pers :: Person{T},
        jsa :: JobSeekersAllowance{T},
        hrs  :: HoursLimits ) :: T where T
        j = zero(T)
        if jsa.abolished
            return j
        end
        if pers.jsa_type == contributory_jsa &&
           pers.usual_hours_worked <= hrs.med &&
           (! has_limited_capactity_for_work( pers ))
            j = pers.age < 25 ? jsa.u25 : jsa.o24
        end
        return j
    end

    """
    Household level calculations for all nmt benefits that don't require knowlege
    of income tax/ni liabilties - not necessarily taxable benefits, just things
    that don't require any kind of net income calculation and so can be done
    before IT/NI.
    """
    function calc_pre_tax_non_means_tested!(
        hhres :: HouseholdResult,
        hh    :: Household,
        sys   :: NonMeansTestedSys,
        hours_limits   :: HoursLimits,
        age_limits :: AgeLimits )
        ## maybe add a benefit unit allocator
        bus = get_benefit_units( hh )
        buno = 1
        for bu in bus
            has_children = size( bu.children )[1] > 0
            bures = hhres.bus[buno]
            for adno in bu.adults
                pers = bu.people[adno]
                pres = bures.pers[adno]
                pres.income[STATE_PENSION] = calc_state_pension(
                    pers,
                    sys.pensions, age_limits );
                pres.income[WIDOWS_PAYMENT] = calc_widows_benefits(
                    pers, has_children, sys.bereavement, sys.widows_pension )
                #
                # FIXME
                # pip/dla can only be claimed
                # if ! reached_state_pension_age( age_limits, pers.age, pers.sex )
                # but claims can run on indefinitely and for now we're just using
                # receipts, so ignore any upper age limits until we model these fully.
                #
                pres.income[sys.dla.care_slot],
                pres.income[sys.dla.mob_slot] = calc_dla( pers, sys.dla );
                pres.income[sys.pip.care_slot],
                pres.income[sys.pip.mob_slot] = calc_pip( pers, sys.pip )

                #
                # .. conversely, this age limit seems safe: a 62 yo female recieving
                # in the data should be disallowed now the pension age has increased.
                #
                if reached_state_pension_age( age_limits, pers.age, pers.sex )
                    pres.income[sys.attendance_allowance.slot] = calc_attendance_allowance( pers, sys.attendance_allowance )
                else
                    pres.income[CONTRIB_EMPLOYMENT_AND_SUPPORT_ALLOWANCE] = calc_esa( pers, sys.esa )
                    pres.income[CONTRIB_JOBSEEKERS_ALLOWANCE] = calc_jsa( pers, sys.jsa, hours_limits )
                    pres.income[MATERNITY_ALLOWANCE] = calc_maternity_allowance( pers, sys.maternity )
                end
                # NON-overlapping rules p1178 go here
            end # ad loop
            buno += 1
        end # bu loop
    end # calc_non_means_tested

    """
    NMT Benefits that require knowlege of tax and NI liabilities and so have to be done
    after a tax calculation - not necessarily tax free as such (CB higher charge). Kinda-sorta
    means-tested bens, I suppose.
    """
    function calc_post_tax_non_means_tested!(
        hhres :: HouseholdResult,
        hh    :: Household,
        sys   :: NonMeansTestedSys,
        age_limits :: AgeLimits )
        ## maybe add a benefit unit allocator
        bus = get_benefit_units( hh )
        buno = 1
        for bu in bus
            has_children = size( bu.children )[1] > 0
            bures = hhres.bus[buno]
            calc_child_benefit!(
                bures,
                bu,
                sys.child_benefit )
            for adno in bu.adults
                pers = bu.people[adno]
                pres = bures.pers[adno]
                pres.income[sys.carers.slot] =
                    calc_carers_allowance( pers, pres, sys.carers )

                if hh.region == Scotland
                    if pres.income[sys.carers.slot] > 0
                        pres.income[SCOTTISH_CARERS_SUPPLEMENT] = sys.carers.scottish_supplement
                    end
                end
                # NON-overlapping rules p1178 go here
            end # ad loop
            buno += 1
        end # bu loop
    end # calc_non_means_tested

end # package non-means-tested


module NastyGenerator

using Random

pn(i) = "P$i"
randvers(k::Int) = VersionNumber(rand(0:k), rand(0:k), rand(0:k))

function randvspec(k::Int)
    lb, ub = randvers(k), randvers(k)
    if lb > ub
        lb, ub = ub, lb
    end

    slb = rand() < 0.1 ? "0" : string(lb)
    sub = rand() < 0.1 ? "*" : string(ub)

    return "$slb-$sub"
end

"""
Generates a random graph with 2 planted solutions (or quasi-solutions if sat==false).
With the right parameters, this is quite hard to solve for large sizes.
For added fun, the two planted solutions are cycles.
We want the solver to find the best of the two, of course.
This is an extremely unrealistic scenario and it's only intended to stress-test the solver.
Note that the "problematic" output assumes that all non-planted versions will be
uninstallable, which is only the case for some regimes of the parameters (e.g. large
enough d).
"""
function generate_nasty(n::Int,             # size of planted solutions
                        m::Int;             # size of the graph
                        k::Int = 10,        # version number limit
                        q::Int = 10,        # versions per package (upper bound)
                        d::Int = 10,        # neighbors per package
                        seed::Integer = 32524,
                        sat::Bool = true    # create a satisfiable problem?
                       )
    @assert m ≥ n
    d ≤ m-1 || @warn "d=$d, should be ≤ m-1=$(m-1)"

    srand(seed)

    allvers = [sort(unique(randvers(k) for j = 1:q)) for i = 1:m]

    planted1 = [rand(2:length(allvers[i])) for i = 1:n]

    planted2 = [rand(1:(planted1[i]-1)) for i = 1:n]

    deps = []
    problematic = []

    # random dependencies
    for i = 1:m, j = 1:length(allvers[i])
        if i ≤ n && (planted1[i] == j || planted2[i] == j)
            if j == planted1[i]
                if i < n
                    push!(deps, [pn(i), allvers[i][j], pn(i+1), "$(allvers[i+1][planted1[i+1]])-*"])
                else
                    if !sat
                        push!(deps, [pn(i), allvers[i][j], pn(1), "0-$(allvers[1][planted2[1]])"])
                    else
                        push!(deps, [pn(i), allvers[i][j], pn(1), "0-*"])
                    end
                end
            else # j == planted2[i]
                if i < n
                    push!(deps, [pn(i), allvers[i][j], pn(i+1), "0-$(allvers[i+1][planted2[i+1]])"])
                else
                    if !sat
                        push!(deps, [pn(i), allvers[i][j], pn(1), "$(allvers[1][planted1[1]])-*"])
                    else
                        push!(deps, [pn(i), allvers[i][j], pn(1), "0-*"])
                    end
                end
            end
            sat || push!(problematic, [pn(i), allvers[i][j]])
            continue
        end

        s = shuffle([1:(i-1); (i+1):m])[1:min(d,m-1)]
        for a in s
            push!(deps, [pn(i), allvers[i][j], pn(a), randvspec(k)])
        end
        push!(problematic, [pn(i), allvers[i][j]])
    end

    reqs = [[pn(1), "*"]]
    # reqs = [[pn(1), string(allvers[1][planted1[1]])]]

    # info("SOLUTION: $([(i,planted1[i]) for i = 1:n])")
    # info("REST: $([(i,length(allvers[i])+1) for i = (n+1):m])")

    want = Dict(pn(i) => allvers[i][planted1[i]] for i = 1:n)

    return deps, reqs, want, problematic
end

end # module


const SVF = MOI.SingleVariable
const VVF = MOI.VectorOfVariables
const SAF{T} = MOI.ScalarAffineFunction{T}
const VAF{T} = MOI.VectorAffineFunction{T}

const VI = MOI.VariableIndex
const CI = MOI.ConstraintIndex

const GT = MOI.GreaterThan
const LT = MOI.LessThan
const ET = MOI.EqualTo

const SCALAR_SETS = Union{
    MOI.GreaterThan{Float64},
    MOI.LessThan{Float64},
    MOI.EqualTo{Float64},
}

const VECTOR_SETS = Union{
    MOI.SecondOrderCone,
    MOI.RotatedSecondOrderCone,
}

# abstract type AsbtractBilevelOptimizer end
# struct SOS1Optimizer{O} <: AsbtractBilevelOptimizer
#     solver::O
#     # options
# end
# function SOS1Optimizer(solver::O) where O
#     return SOS1Optimizer{O}(solver)
# end

struct Complement#{M1 <: MOI.ModelLike, M2 <: MOI.ModelLike, F, S}
    is_vec
    # primal::M1
    func_w_cte#::F
    set_w_zero#::S
    # dual::M2
    variable#::VI
    # var_set#::S2
end

abstract type BilevelSolverMode{T} end

mutable struct SOS1Mode{T} <: BilevelSolverMode{T}
    epsilon::T
    function SOS1Mode()
        return new{Float64}(zero(Float64))
    end
end

mutable struct PositiveSOS1Mode{T} <: BilevelSolverMode{T}
    epsilon::T
    function PositiveSOS1Mode()
        return new{Float64}(zero(Float64))
    end
end

mutable struct ComplementMode{T} <: BilevelSolverMode{T}
    epsilon::T
    function ComplementMode()
        return new{Float64}(zero(Float64))
    end
end

mutable struct ComplementWithSlackMode{T} <: BilevelSolverMode{T}
    epsilon::T
    function ComplementWithSlackMode()
        return new{Float64}(zero(Float64))
    end
end

mutable struct ProductMode{T} <: BilevelSolverMode{T}
    epsilon::T
    function ProductMode()
        return new{Float64}(zero(Float64))
    end
    function ProductMode{T}(eps::T) where T
        return new{Float64}(eps)
    end
end

mutable struct ProductWithSlackMode{T} <: BilevelSolverMode{T}
    epsilon::T
    function ProductWithSlackMode()
        return new{Float64}(zero(Float64))
    end
    function ProductWithSlackMode{T}(eps::T) where T
        return new{Float64}(eps)
    end
end

@enum IndicatorSetting ZERO_ONE ZERO_ZERO ONE_ONE

mutable struct IndicatorMode{T} <: BilevelSolverMode{T}
    mode::IndicatorSetting
    function IndicatorMode()
        return new{Float64}(ONE_ONE)
    end
end

abstract type StrongDualityMode{T} <: BilevelSolverMode{T} end

mutable struct StrongDualityEqualityMode{T} <: StrongDualityMode{T}
    function StrongDualityEqualityMode()
        return new{Float64}()
    end
end
mutable struct StrongDualityInequalityMode{T} <: StrongDualityMode{T}
    epsilon::T
    function StrongDualityInequalityMode()
        return new{Float64}(zero(Float64))
    end
    function StrongDualityInequalityMode{T}(eps::T) where T
        return new{Float64}(eps)
    end
end

ignore_dual_objective(::BilevelSolverMode{T}) where T = true
ignore_dual_objective(::StrongDualityMode{T}) where T = false

function accept_vector_set(mode::BilevelSolverMode{T}, con::Complement) where T
    if con.is_vec
        error("Set $(typeof(con.set_w_zero)) is not accepted when solution method is $(typeof(mode))")
    end
    return nothing
end
accept_vector_set(::ProductMode{T}, ::Complement) where T = nothing
accept_vector_set(::ProductWithSlackMode{T}, ::Complement) where T = nothing

function get_canonical_complements(primal_model, primal_dual_map)
    map = primal_dual_map.primal_con_dual_var
    out = Complement[]
    for ci in keys(map)
        con = get_canonical_complement(primal_model, map, ci)
        push!(out, con)
    end
    return out
end
function get_canonical_complement(primal_model, map,
    ci::CI{F,S}) where {F, S<:VECTOR_SETS}
    T = Float64
    func = MOI.get(primal_model, MOI.ConstraintFunction(), ci)::F
    set = MOI.get(primal_model, MOI.ConstraintSet(), ci)::S
    dim = MOI.dimension(set)
    # vector sets have no constant
    # for i in 1:dim
    #     func.constant[i] = Dualization.set_dot(i, set, T) *
    #         Dualization.get_scalar_term(primal_model, i, ci)
    # end
    # todo - set dot on function
    con = Complement(true, func, set_with_zero(set), map[ci])
    return con
end
function get_canonical_complement(primal_model, map,
    ci::CI{F,S}) where {F, S<:SCALAR_SETS}
    T = Float64
    func = MOI.get(primal_model, MOI.ConstraintFunction(), ci)::F
    set = MOI.get(primal_model, MOI.ConstraintSet(), ci)::S
    constant = Dualization.set_dot(1, set, T) *
        Dualization.get_scalar_term(primal_model, 1, ci)
    if F == MOI.SingleVariable
        func = MOIU.operate(+, T, func, constant)
    else
        func.constant = constant
    end
    # todo - set dot on function
    con = Complement(false, func, set_with_zero(set), map[ci][1])
    return con
end

function set_with_zero(set::S) where {S<:SCALAR_SETS} where T
    return S(0.0)
end
function set_with_zero(set)
    return copy(set)
end

function build_bilevel(
    upper::MOI.ModelLike, lower::MOI.ModelLike,
    link::Dict{VI,VI}, upper_variables::Vector{VI},
    mode,
    upper_var_lower_ctr::Dict{VI,CI} = Dict{VI,CI}()
    )

    # Start with an empty problem
    moi_mode = MOIU.AUTOMATIC
    # m = MOIU.CachingOptimizer(MOIU.Model{Float64}(), moi_mode)
    m = MOIU.CachingOptimizer(MOIU.UniversalFallback(MOIU.Model{Float64}()), moi_mode)

    # add the first level
    copy_names = false

    #=
        Initialize Lower DUAL level model
    =#
    # dualize the second level
    dual_problem = dualize(lower,
        dual_names = DualNames("dual_","dual_"),
        variable_parameters = upper_variables,
        ignore_objective = ignore_dual_objective(mode))
    lower_dual = dual_problem.dual_model
    lower_primal_dual_map = dual_problem.primal_dual_map

    #=
        Pass Upper level model
    =#

    # key are from src, value are from dest
    upper_idxmap = MOIU.default_copy_to(m, upper, copy_names)
    pass_names(m, upper, upper_idxmap)

    #=
        Pass Lower level model
    =#

    # cache and delete lower objective
    if !ignore_dual_objective(mode)
        # get primal obj
        tp_primal_obj = MOI.get(lower, MOI.ObjectiveFunctionType())
        @assert tp_primal_obj !== nothing
        lower_primal_obj = MOI.get(lower, MOI.ObjectiveFunction{tp_primal_obj}())
        # delete dual obj
        # MOI.set(lower, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{Float64}[], 0.0))
    end

    # initialize map to lower level model
    lower_idxmap = MOIU.IndexMap()
    for (upper_key, lower_val) in link
        lower_idxmap[lower_val] = upper_idxmap[upper_key]
    end

    # append the second level primal
    append_to(m, lower, lower_idxmap, copy_names, allow_single_bounds = true)
    pass_names(m, lower, lower_idxmap)

    #=
        Pass Dual of Lower level model
    =#

    # initialize map to lower level model
    if !ignore_dual_objective(mode)
        # get dual obj
        tp_dual_obj = MOI.get(lower_dual, MOI.ObjectiveFunctionType())
        @assert tp_dual_obj !== nothing
        lower_dual_obj = MOI.get(lower_dual, MOI.ObjectiveFunction{tp_dual_obj}())
        # delete dual obj
        # MOI.set(lower_dual, MOI.ObjectiveFunction{MOI.ScalarAffineFunction{Float64}}(), MOI.ScalarAffineFunction(MOI.ScalarAffineTerm{Float64}[], 0.0))
    end

    # initialize map to lower level model
    lower_dual_idxmap = MOIU.IndexMap()
    # for QP's there are dual variable that are tied to:
    # primal variables
    for (lower_primal_var_key, lower_dual_quad_slack_val) in lower_primal_dual_map.primal_var_dual_quad_slack
        lower_dual_idxmap[lower_dual_quad_slack_val] = lower_idxmap[lower_primal_var_key]
    end
    # and to upper level variable which are lower level parameters
    for (lower_primal_param_key, lower_dual_param_val) in lower_primal_dual_map.primal_parameter
        lower_dual_idxmap[lower_dual_param_val] = lower_idxmap[lower_primal_param_key]
    end
    # Dual variables might appear in the upper level
    for (upper_var, lower_con) in upper_var_lower_ctr
        var = lower_primal_dual_map.primal_con_dual_var[lower_con][1] # TODO check this scalar
        lower_dual_idxmap[var] = upper_idxmap[upper_var]
    end

    # append the second level dual
    append_to(m, lower_dual, lower_dual_idxmap, copy_names)
    pass_names(m, lower_dual, lower_dual_idxmap)

    #=
        Additional Optimiality conditions (to complete the KKT)
    =#

    if ignore_dual_objective(mode)
        # complementary slackness
        comps = get_canonical_complements(lower, lower_primal_dual_map)
        for comp in comps
            if !is_equality(comp.set_w_zero)
                accept_vector_set(mode, comp)
                add_complement(mode, m, comp, lower_idxmap, lower_dual_idxmap)
            else
                # println("eq in complement")
            end
        end
    else
        # strong duality
        lower_dual_obj
        lower_primal_obj
        add_strong_duality(mode, m, lower_primal_obj, lower_dual_obj, lower_idxmap, lower_dual_idxmap)
    end

    return m, upper_idxmap, lower_idxmap, lower_primal_dual_map, lower_dual_idxmap
end

function add_strong_duality(mode::StrongDualityEqualityMode{T}, m, primal_obj, dual_obj,
    idxmap_primal, idxmap_dual) where T

    primal = MOIU.map_indices.(Ref(idxmap_primal), primal_obj)
    dual   = MOIU.map_indices.(Ref(idxmap_dual), dual_obj)

    func = MOIU.operate(-, T, primal, dual)

    c = MOI.add_constraint(m, func, MOI.EqualTo(zero(T)))

    MOI.set(m, MOI.ConstraintName(), c, "lower_strong_duality")

    return c
end
function add_strong_duality(mode::StrongDualityInequalityMode{T}, m, primal_obj, dual_obj,
    idxmap_primal, idxmap_dual) where T

    primal = MOIU.map_indices.(Ref(idxmap_primal), primal_obj)
    dual   = MOIU.map_indices.(Ref(idxmap_dual), dual_obj)

    func = MOIU.operate(-, T, primal, dual)

    func_up = MOIU.operate(-, T, func, mode.epsilon)
    c_up = MOI.add_constraint(m, func_up, MOI.LessThan(zero(T)))
    MOI.set(m, MOI.ConstraintName(), c_up, "lower_strong_duality_up")

    func_lo = MOIU.operate(+, T, func, mode.epsilon)
    c_lo = MOI.add_constraint(m, func_lo, MOI.GreaterThan(zero(T)))
    MOI.set(m, MOI.ConstraintName(), c_lo, "lower_strong_duality_lo")

    return c_up, c_lo
end

function add_complement(mode::ComplementMode{T}, m, comp::Complement, idxmap_primal, idxmap_dual) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    f_dest = MOIU.map_indices.(Ref(idxmap_primal), f)

    dual = idxmap_dual[v]

    new_f = MOIU.operate(vcat, T, f_dest, MOI.SingleVariable(dual))

    c1 = MOI.add_constraint(m,
        new_f,
        MOI.Complements(1))

    nm = MOI.get(m, MOI.VariableName(), dual)
    MOI.set(m, MOI.ConstraintName(), c1, "compl_compl_($(nm))")

    return c1
end

function add_complement(mode::ComplementWithSlackMode{T}, m, comp::Complement, idxmap_primal, idxmap_dual) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    slack, slack_in_set = MOI.add_constrained_variable(m, s)
    f_dest = MOIU.map_indices.(Ref(idxmap_primal), f)
    new_f = MOIU.operate(-, T, f_dest, MOI.SingleVariable(slack))
    equality = MOIU.normalize_and_add_constraint(m, new_f, MOI.EqualTo(zero(T)))

    # add start value to slack
    val = MOIU.eval_variables(
        x-> nothing_to_nan(MOI.get(m, MOI.VariablePrimalStart(), x)), f_dest)
    if !isnan(val)
        MOI.set(m, MOI.VariablePrimalStart(), slack, val)
    end

    dual = idxmap_dual[v]
    c1 = MOI.add_constraint(m,
        MOI.VectorOfVariables([slack, dual]),
        MOI.Complements(1))

    nm = MOI.get(m, MOI.VariableName(), dual)
    MOI.set(m, MOI.VariableName(), slack, "slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), slack_in_set, "bound_slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), equality, "eq_slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), c1, "compl_complWslk_($(nm))")

    return slack, slack_in_set, equality, c1
end

function add_complement(mode::SOS1Mode{T}, m, comp::Complement, idxmap_primal, idxmap_dual) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    slack, slack_in_set = MOI.add_constrained_variable(m, s)
    f_dest = MOIU.map_indices.(Ref(idxmap_primal), f)
    new_f = MOIU.operate(-, T, f_dest, MOI.SingleVariable(slack))
    equality = MOIU.normalize_and_add_constraint(m, new_f, MOI.EqualTo(zero(T)))

    dual = idxmap_dual[v]
    c1 = MOI.add_constraint(m,
        MOI.VectorOfVariables([slack, dual]),
        MOI.SOS1([1.0, 2.0]))

    nm = MOI.get(m, MOI.VariableName(), dual)
    MOI.set(m, MOI.VariableName(), slack, "slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), slack_in_set, "bound_slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), equality, "eq_slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), c1, "compl_sos1_($(nm))")

    return slack, slack_in_set, equality, c1
end

function add_complement(mode::PositiveSOS1Mode{T}, m, comp::Complement, idxmap_primal, idxmap_dual) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    f_dest = MOIU.map_indices.(Ref(idxmap_primal), f)

    if typeof(s) <: MOI.LessThan # 0
        # requires flipping
        # flipped slack
        slack, slack_in_set = MOI.add_constrained_variable(m, MOI.GreaterThan{T}(0.0))
        new_f = MOIU.operate(+, T, f_dest, MOI.SingleVariable(slack))
        # flipped dual
        real_dual = idxmap_dual[v]
        dual, dual_in_set = MOI.add_constrained_variable(m, MOI.GreaterThan{T}(0.0))
        # dual + real_dual == 0
        opposite = MOIU.normalize_and_add_constraint(m,
            MOI.ScalarAffineFunction(
                [MOI.ScalarAffineTerm(1.0, real_dual),
                 MOI.ScalarAffineTerm(1.0, dual)],
                0.0),
            MOI.EqualTo(zero(T)))

        nm = MOI.get(m, MOI.VariableName(), dual)
        MOI.set(m, MOI.VariableName(), slack, "flip_dual_($(nm))")
        MOI.set(m, MOI.ConstraintName(), slack_in_set, "flip_dual_in_set_($(nm))")
        MOI.set(m, MOI.ConstraintName(), opposite, "flip_dual_eq_($(nm))")
    elseif typeof(s) <: MOI.GreaterThan # 0
        slack, slack_in_set = MOI.add_constrained_variable(m, s)
        new_f = MOIU.operate(-, T, f_dest, MOI.SingleVariable(slack))
        dual = idxmap_dual[v]
    else
        error("Unexpected set type: $s, while building complment constraints.")
    end

    equality = MOIU.normalize_and_add_constraint(m, new_f, MOI.EqualTo(zero(T)))

    c1 = MOI.add_constraint(m,
        MOI.VectorOfVariables([slack, dual]),
        MOI.SOS1([1.0, 2.0]))

    nm = MOI.get(m, MOI.VariableName(), dual)
    MOI.set(m, MOI.VariableName(), slack, "slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), slack_in_set, "bound_slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), equality, "eq_slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), c1, "compl_sos1_($(nm))")

    return slack, slack_in_set, equality, c1
end

function flip_set(set::MOI.LessThan{T}) where T
    return MOI.GreaterThan{T}(0.0)
end
function flip_set(set::MOI.GreaterThan{T}) where T
    return MOI.LessThan{T}(0.0)
end
is_equality(set::S) where {S<:MOI.AbstractSet} = false
is_equality(set::MOI.EqualTo{T}) where T = true
is_equality(set::MOI.Zeros) = true

only_variable_functions(v::MOI.VariableIndex) = MOI.SingleVariable(v)
only_variable_functions(v::Vector{MOI.VariableIndex}) = MOI.VectorOfVariables(v)

function add_complement(mode::ProductMode{T}, m, comp::Complement, idxmap_primal, idxmap_dual) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    eps = mode.epsilon

    f_dest = MOIU.map_indices(x->idxmap_primal[x], f)

    dual = comp.is_vec ? map(x->idxmap_dual[x], v) : idxmap_dual[v]

    new_f = MOIU.operate(dot, T, f_dest, only_variable_functions(dual))
    new_f1 = MOIU.operate(-, T, new_f, eps)
    c1 = MOI.add_constraint(m,
        new_f1,
        MOI.LessThan{T}(0.0))
    if comp.is_vec # conic
        new_f2 = MOIU.operate(+, T, new_f, eps)
        c2 = MOI.add_constraint(m,
            new_f2,
            MOI.GreaterThan{T}(0.0))
    end

    # c1 = MOI.add_constraint(m,
    #     new_f,
    #     MOI.EqualTo(zero(T)))

    nm = comp.is_vec ? MOI.get.(m, MOI.VariableName(), dual) : MOI.get(m, MOI.VariableName(), dual)

    MOI.set(m, MOI.ConstraintName(), c1, "compl_prod_($(nm))")
    if comp.is_vec
        MOI.set(m, MOI.ConstraintName(), c2, "compl_prod2_($(nm))")
    end

    return c1
end

nothing_to_nan(val) = ifelse(val === nothing, NaN, val)

function add_complement(mode::ProductWithSlackMode{T}, m, comp::Complement, idxmap_primal, idxmap_dual) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    eps = mode.epsilon

    slack, slack_in_set = MOI.add_constrained_variable(m, s)
    f_dest = MOIU.map_indices.(Ref(idxmap_primal), f)
    new_f = MOIU.operate(-, T, f_dest, MOI.SingleVariable(slack))
    equality = MOIU.normalize_and_add_constraint(m, new_f, MOI.EqualTo(zero(T)))

    # add start value to slack
    val = MOIU.eval_variables(
        x-> nothing_to_nan(MOI.get(m, MOI.VariablePrimalStart(), x)), f_dest)
    if !isnan(val)
        MOI.set(m, MOI.VariablePrimalStart(), slack, val)
    end

    dual = idxmap_dual[v]

    prod_f = MOIU.operate(*, T, MOI.SingleVariable(slack), MOI.SingleVariable(dual))

    prod_f1 = MOIU.operate(-, T, prod_f, eps)
    c1 = MOI.add_constraint(m,
        prod_f1,
        MOI.LessThan{Float64}(0.0))
    if comp.is_vec # conic
        prod_f2 = MOIU.operate(+, T, prod_f, eps)
        c2 = MOI.add_constraint(m,
            prod_f2,
            MOI.GreaterThan{Float64}(0.0))
        MOI.set(m, MOI.ConstraintName(), c1, "compl_prodWslk2_($(nm))")
    end

    # c1 = MOI.add_constraint(m,
    #     prod_f1,
    #     MOI.EqualTo(zero(T)))

    # TODO dont add slack if single variable
    # TODO use dot product

    nm = MOI.get(m, MOI.VariableName(), dual)
    MOI.set(m, MOI.VariableName(), slack, "slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), slack_in_set, "bound_slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), equality, "eq_slk_($(nm))")
    MOI.set(m, MOI.ConstraintName(), c1, "compl_prodWslk_($(nm))")

    return slack, slack_in_set, equality, c1
end

function add_complement(mode::IndicatorMode{T}, m, comp::Complement, idxmap_primal, idxmap_dual) where T
    f = comp.func_w_cte
    s = comp.set_w_zero
    v = comp.variable

    method = mode.mode

    f_dest = MOIU.map_indices(x->idxmap_primal[x], f)

    dual = idxmap_dual[v]

    if comp.is_vec
        @show s
        error("vector constraint not supported")
    end

    nm = comp.is_vec ? MOI.get.(m, MOI.VariableName(), dual) : MOI.get(m, MOI.VariableName(), dual)

    vb1 = MOI.add_variable(m)
    MOI.set(m, MOI.VariableName(), vb1, "compl_bin1_($(nm))")

    cb1 = MOI.add_constraint(m, SVF(vb1), MOI.ZeroOne())
    if method == ONE_ONE || method == ZERO_ZERO
        # second binary
        vb2 = MOI.add_variable(m)
        cb2 = MOI.add_constraint(m, SVF(vb2), MOI.ZeroOne())
        MOI.set(m, MOI.VariableName(), vb2, "compl_bin2_($(nm))")

        # z1 + z2 == 1
        fb = MOIU.operate(+, T, SVF(vb1), SVF(vb2))
        cb = MOI.add_constraint(m, fb, MOI.EqualTo{T}(one(T)))
        MOI.set(m, MOI.ConstraintName(), cb, "compl_sum_bin_($(nm))")
    else
        vb2 = vb1
    end

    f1 = MOIU.operate(vcat, T, SVF(vb1), f_dest)
    f2 = MOIU.operate(vcat, T, SVF(vb2), SVF(dual))

    if method == ONE_ONE
        s1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.EqualTo(zero(T)))
        s2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.EqualTo(zero(T)))
    elseif method == ZERO_ZERO
        s1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}(MOI.EqualTo(zero(T)))
        s2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}(MOI.EqualTo(zero(T)))
    else
        s1 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ONE}(MOI.EqualTo(zero(T)))
        s2 = MOI.IndicatorSet{MOI.ACTIVATE_ON_ZERO}(MOI.EqualTo(zero(T)))
    end

    c1 = MOI.add_constraint(m, to_vector_affine(f1), s1)
    c2 = MOI.add_constraint(m, to_vector_affine(f2), s2)

    MOI.set(m, MOI.ConstraintName(), c1, "compl_ind1_($(nm))")
    MOI.set(m, MOI.ConstraintName(), c2, "compl_ind2_($(nm))")

    return c1
end

function to_vector_affine(f::MOI.VectorAffineFunction{T}) where T
    return f
end
function to_vector_affine(f::MOI.VectorOfVariables)
    return MOI.VectorAffineFunction{Float64}(f)
end

function pass_names(dest, src, map)
    for vi in MOI.get(src, MOI.ListOfVariableIndices())
        name = MOI.get(src, MOI.VariableName(), vi)
        if name != ""
            MOI.set(dest, MOI.VariableName(), map[vi], name)
        end
    end
    for (F,S) in MOI.get(src, MOI.ListOfConstraints())
        for con in MOI.get(src, MOI.ListOfConstraintIndices{F,S}())
            name = MOI.get(src, MOI.ConstraintName(), con)
            if name != ""
                MOI.set(dest, MOI.ConstraintName(), map[con], name)
            end
        end
    end
end

function append_to(dest::MOI.ModelLike, src::MOI.ModelLike, idxmap, copy_names::Bool; allow_single_bounds::Bool = true)
    # MOI.empty!(dest)

    # idxmap = MOIU.IndexMap()

    vis_src = MOI.get(src, MOI.ListOfVariableIndices())
    constraint_types = MOI.get(src, MOI.ListOfConstraints())
    single_variable_types = [S for (F, S) in constraint_types
                             if F == MOI.SingleVariable && allow_single_bounds]
    vector_of_variables_types = [S for (F, S) in constraint_types
                                 if F == MOI.VectorOfVariables]

    # The `NLPBlock` assumes that the order of variables does not change (#849)
    if MOI.NLPBlock() in MOI.get(src, MOI.ListOfModelAttributesSet())
        error("not nlp for now")
        vector_of_variables_not_added = [
            MOI.get(src, MOI.ListOfConstraintIndices{MOI.VectorOfVariables, S}())
            for S in vector_of_variables_types
        ]
        single_variable_not_added = [
            MOI.get(src, MOI.ListOfConstraintIndices{MOI.SingleVariable, S}())
            for S in single_variable_types
        ]
    else
        vector_of_variables_not_added = [
            MOIU.copy_vector_of_variables(dest, src, idxmap, S)
            for S in vector_of_variables_types
        ]
        single_variable_not_added = [
            MOIU.copy_single_variable(dest, src, idxmap, S)
            for S in single_variable_types
        ]
    end

    # MOIU.copy_free_variables(dest, idxmap, vis_src, MOI.add_variables)
    for vi in vis_src
        if !haskey(idxmap.varmap, vi)
            var = MOI.add_variable(dest)
            idxmap.varmap[vi] = var
        end
    end

    # Copy variable attributes
    MOIU.pass_attributes(dest, src, copy_names, idxmap, vis_src)

    # Copy model attributes
    # attention HERE to no pass objective functions!
    # pass_attributes(dest, src, copy_names, idxmap)

    # Copy constraints
    MOIU.pass_constraints(dest, src, copy_names, idxmap,
                     single_variable_types, single_variable_not_added,
                     vector_of_variables_types, vector_of_variables_not_added)

    return idxmap
end

using LinearAlgebra

# scalar
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction{T}}},
    ::Type{T}
    ) where T
    MOI.ScalarAffineFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{T},
    ::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction{T}}}
    ) where T
    MOI.ScalarAffineFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction{T}}},
    ::Type{<:Union{MOI.SingleVariable, MOI.ScalarAffineFunction{T}}}
    ) where T
    MOI.ScalarQuadraticFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{MOI.ScalarQuadraticFunction{T}},
    ::Type{T}
    ) where T
    MOI.ScalarQuadraticFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{T},
    ::Type{MOI.ScalarQuadraticFunction{T}}
    ) where T
    MOI.ScalarQuadraticFunction{T}
end
# flip
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        MOI.SingleVariable,
        MOI.ScalarAffineFunction{T},
        MOI.ScalarQuadraticFunction{T}
        },
    α::T) where T
    return MOIU.operate(LinearAlgebra.dot, T, α, f)
end
# pass to *
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        T,
        MOI.SingleVariable,
        MOI.ScalarAffineFunction{T}
        },
    g::Union{
        MOI.SingleVariable,
        MOI.ScalarAffineFunction{T}
        }
    ) where T
    return MOIU.operate(*, T, f, g)
end
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    α::T,
    f::MOI.ScalarQuadraticFunction{T}
    ) where T
    return MOIU.operate(*, T, f, α)
end

# vector
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}},
    ::Type{Vector{T}}
    ) where T
    MOI.VectorAffineFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{Vector{T}},
    ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}}
    ) where T
    MOI.VectorAffineFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}},
    ::Type{<:Union{MOI.VectorOfVariables, MOI.VectorAffineFunction{T}}}
    ) where T
    MOI.VectorQuadraticFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{MOI.VectorQuadraticFunction{T}},
    ::Type{Vector{T}}
    ) where T
    MOI.VectorQuadraticFunction{T}
end
function MOIU.promote_operation(::typeof(LinearAlgebra.dot), ::Type{T},
    ::Type{Vector{T}},
    ::Type{MOI.VectorQuadraticFunction{T}}
    ) where T
    MOI.VectorQuadraticFunction{T}
end
# flip
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T},
        MOI.VectorQuadraticFunction{T}
        },
    α::Vector{T}) where T
    return MOIU.operate(LinearAlgebra.dot, T, α, f)
end
# pass to _operate(LinearAlgebra.dot, ...)
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        Vector{T},
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T}
        },
    g::Union{
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T}
        }
    ) where T
    return _operate(LinearAlgebra.dot, T, f, g)
end
function MOIU.operate(::typeof(LinearAlgebra.dot), ::Type{T},
    α::T,
    f::MOI.VectorQuadraticFunction{T}
    ) where T
    return _operate(LinearAlgebra.dot, T, f, α)
end
function _operate(::typeof(LinearAlgebra.dot), ::Type{T},
    f::Union{
        Vector{T},
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T},
        MOI.VectorQuadraticFunction{T}
        },
    g::Union{
        MOI.VectorOfVariables,
        MOI.VectorAffineFunction{T},
        MOI.VectorQuadraticFunction{T}
    }) where T

    dim = MOI.output_dimension(g)
    if MOI.output_dimension(f) != dim
        throw(DimensionMismatch("f and g are of different MOI.output_dimension's!"))
    end

    fs = MOIU.scalarize(f)
    gs = MOIU.scalarize(g)

    out = MOIU.operate(*, T, fs[1], gs[1])
    for i in 2:dim
        MOIU.operate!(+, T, out, MOIU.operate(*, T, fs[i], gs[i]))
    end

    return out
end
MOIU.scalarize(v::Vector{T}) where T<:Number = v
MOI.output_dimension(v::Vector{T}) where T<:Number = length(v)


# Find and add MDP's path to Julia search path
cdir = pwd(); ii = findlast("MDP", cdir);
if ii===nothing
    # MDPpath = /path/to/MDP;
    display("MDP's path is not found. Please uncomment the above line and set the path to MDP.");
else
    MDPpath = cdir[1:ii[end]] * "/";
    include(MDPpath * "Installation/setpath.jl");
end

# External packages
using Revise, DelimitedFiles, SymPy

# MDP packages
using Preprocessing, Gencode, Postprocessing

app = Preprocessing.initializeapp(cdir,MDPpath);
#app.buildexec = 1;
app.platform = "cpu";    # change it to app.platform = "gpu" => MDP runs on Nvidia GPU
app.unitstyle = "metal"; # unit system
app.pbc = [1 1 1];       # periodic boundary conditions
nlattice = 4;            # note that # of atoms = 2 * nlattice^3
app.appname = "snapInP" * string(nlattice);

#  neighbor list
app.neighskin = 1.0;
app.neighcell = 0;       # 0 -> O(N^2) algorithm, 1 -> Cell-linked list algorithm to form neighbor list
app.neighmax = 150;      # maximum number of neighbors allowed
app.neighevery = 1;      # perform neighbor rebuild list for "every" iterations
app.neighdelay = 0;      # delay neighbor rebuild
app.neighcheck = 1;      # 0 -> neighbor rebuild by delay and every settings,
                         # 1 -> neighbor rebuild occurs if some atom has moved more than half the skin distance
app.neighpair = 0;       # 0 -> full neighbor list, 1 -> half neighbor list for pair potentials

# atom
app.atommasses = [114.76 30.98];
app.atomcharges = [0 0];

# lattice, region, domain
basisatomtype = [1 1 1 1 2 2 2 2]; # 1 => In, 2 => P
app.lattice = setlattice("diamond", 5.83,[1.0 0.0 0.0],[0.0 1.0 0.0],[0.0 0.0 1.0],[1.0 1.0 1.0],basisatomtype);
app.region = setregion([nlattice nlattice nlattice]);

#app.lattice = Preprocessing.setlattice("bcc", 3.316);

# nonparametric potential descriptors
app.descriptor = "snap";    # (snap, shp)
app.snaprcutfac = 1.0;
app.snaptwojmax = 6;
app.snaprfac0 = 0.99363;
app.snaprmin0 = 0;
app.snapbzeroflag = 1;
app.snapquadraticflag = 0;
app.snapswitchflag = 1;
app.snapchemflag = 1;
app.snapbnormflag = 1;
app.snapwselfallflag = 1;
app.snapnelem = 2;    # # of elements
app.snapncoeff = 241;  # # of coeffients per element
app.snapelemradius = [3.81205 3.82945]; # radius per element
app.snapelemweight = [1 0.929316];   # weight per element
app.snapcoefffile = "snapInP.coeff";
app.rcutml = 2*app.snaprcutfac*maximum(app.snapelemradius[:]); # maximum cut-off radius for nonparametric potential

# parametric potential descriptors
app.potentialfile = "ZBLpotentialInP.jl";   # empirical potential file
include(app.potentialfile);  # include the potential file
zblcutinner = 4;
zblcutouter = 4.2;
zblz1 = 49;
zblz2 = 15;
app.pot2a = reshape([1],1,1);     # list of nonbonded pair potentials
app.mu2a = [zblcutinner zblcutouter zblz1 zblz2];    # parameters for ZBL potential
app.rcut2a = reshape([zblcutouter],1,1);  # cut-off radius

# time integration and outputs
app.ensemblemode = "nve";
app.time = 0.0;             # initial time
app.dt = 0.5e-3;            # time step size
app.ntimesteps = 100;       # # time steps
app.globalfreq = 10;

temp = 300;  seed0 = 4928459; distflag = 0; sumflag = 0; loopflag = 2; momflag = 1; rotflag = 0;
app.createvelocity = [temp seed0 distflag sumflag loopflag momflag rotflag];

# generate input files and run simulation
app = mdp(app);


module SIMDAlignment

export
    # wrapper types
    seq_t,
    submat_t,
    alignment_t,
    # functions
    paralign_score

import Bio
using Bio.Seq
using Bio.Align

# 1-byte alphabets
typealias Byte Union{Int8,UInt8,DNANucleotide,RNANucleotide,AminoAcid}

# sequence
immutable seq_t
    data::Ptr{UInt8}
    len::Csize_t
    offset::Csize_t
    reversed::Bool
    packed::Bool
end

function Base.convert{T<:Byte}(::Type{seq_t}, seq::Vector{T})
    return seq_t(pointer(seq), length(seq), 0, false, false)
end

function Base.call(::Type{seq_t}, seq::Vector, offset::Int=0, reversed::Bool=false)
    return seq_t(pointer(seq), length(seq), offset, reversed, false)
end

function Base.convert(::Type{seq_t}, seq::NucleotideSequence)
    byteseq = reinterpret(UInt8, seq.data)
    return seq_t(pointer(byteseq), length(seq), seq.part.start - 1, false, true)
end

function Base.call{T}(::Type{seq_t}, seq::NucleotideSequence{T}, reversed::Bool=false)
    byteseq = reinterpret(UInt8, seq.data)
    len = length(seq)
    if len == 0
        offset = 0
    else
        offset = (reversed ? seq.part.stop : seq.part.start) - 1
    end
    return seq_t(pointer(byteseq), length(seq), offset, reversed, true)
end

# substitution matrix
immutable submat_t{T}
    data::Ptr{T}
    size::Cint
end

function Base.convert{T}(::Type{submat_t}, submat::Matrix{T})
    m, n = size(submat)
    @assert m == n
    return submat_t{T}(pointer(submat), convert(Cint, n))
end

function Base.convert{T}(::Type{submat_t}, submat::SubstitutionMatrix{T})
    return submat_t(submat.data)
end

# alignment result
type alignment_t
    score::Int64
    trace::Ptr{UInt8}
    seqlen::Csize_t
    reflen::Csize_t
    endpos_seq::Csize_t
    endpos_ref::Csize_t
    function alignment_t()
        return new(0, C_NULL, 0, 0, 0, 0)
    end
end

function Bio.Align.score(aln::alignment_t)
    return aln.score
end

function Base.show(io::IO, aln::alignment_t)
    print(io, aln.score)
end

const libsimdalign = Pkg.dir("SIMDAlignment", "deps", "libsimdalign.so")

function make_buffer()
    buffer = ccall((:make_buffer, libsimdalign), Ptr{Void}, ())
    @assert buffer != C_NULL "failed to allocate a buffer"
    return buffer
end

function free_buffer(buffer)
    ccall((:free_buffer, libsimdalign), Void, (Ptr{Void},), buffer)
end

@generated function paralign_score{score_t}(submat::Matrix{score_t}, gap_open::score_t, gap_extend::score_t, seq::seq_t, refs::Vector{seq_t})
    func = score_t === Int8  ? :(:paralign_score_i8x32)  :
           score_t === Int16 ? :(:paralign_score_i16x16) :
           score_t === Int32 ? :(:paralign_score_i32x8) :
           error("not supported type: $score_t")
    quote
        alns = Vector{alignment_t}()
        for _ in 1:length(refs)
            push!(alns, alignment_t())
        end
        buffer = make_buffer()
        ret = ccall(
            ($(func), libsimdalign),
            Cint,
            (Ptr{Void}, submat_t{score_t}, score_t, score_t, seq_t, Ptr{seq_t}, Cint, Ptr{Void}),
            buffer, submat_t(submat), gap_open, gap_extend, seq, pointer(refs), length(refs), alns
        )
        free_buffer(buffer)
        @assert ret == 0 "failed to align"
        return alns
    end
end

function paralign_score{score_t}(submat::Union{Matrix{score_t},SubstitutionMatrix{score_t}}, gap_open, gap_extend, seq, refs)
    paralign_score(
        convert(Matrix{score_t}, submat),
        score_t(gap_open),
        score_t(gap_extend),
        seq_t(seq),
        [seq_t(ref) for ref in refs]
    )
end

end # module


const Optional{T} = Union{Nothing, T}

struct AxisAspect
    aspect::Float32
end

struct DataAspect end

"""
LinearTicks with ideally a number of `n_ideal` tick marks.
"""
struct LinearTicks
    n_ideal::Int

    function LinearTicks(n_ideal)
        if n_ideal <= 0
            error("Ideal number of ticks can't be smaller than 0, but is $n_ideal")
        end
        new(n_ideal)
    end
end

struct WilkinsonTicks
    k_ideal::Int
    k_min::Int
    k_max::Int
    Q::Vector{Tuple{Float64, Float64}}
    granularity_weight::Float64
    simplicity_weight::Float64
    coverage_weight::Float64
    niceness_weight::Float64
    min_px_dist::Float64
end

"""
Like LinearTicks but for multiples of `multiple`.
Example where approximately 5 numbers should be found
that are multiples of pi, printed like "1π", "2π", etc.:

```
MultiplesTicks(5, pi, "π")
```
"""
struct MultiplesTicks
    n_ideal::Int
    multiple::Float64
    suffix::String
end

mutable struct LineAxis
    parent::Scene
    protrusion::Node{Float32}
    attributes::Attributes
    elements::Dict{Symbol, Any}
    tickpositions::Node{Vector{Point2f0}}
    tickvalues::Node{Vector{Float32}}
    ticklabels::Node{Vector{String}}
end

struct LimitReset end

mutable struct RectangleZoom
    active::Bool
    restrict_x::Bool
    restrict_y::Bool
    from::Union{Nothing, Point2f0}
    to::Union{Nothing, Point2f0}
    rectnode::Observable{FRect2D}
    plots::Vector{AbstractPlot}
end

struct ScrollZoom
    speed::Float32
    reset_timer::Ref{Any}
    prev_xticklabelspace::Ref{Any}
    prev_yticklabelspace::Ref{Any}
    reset_delay::Float32
end

struct DragPan
    reset_timer::Ref{Any}
    prev_xticklabelspace::Ref{Any}
    prev_yticklabelspace::Ref{Any}
    reset_delay::Float32
end

struct ScrollEvent
    x::Float32
    y::Float32
end

struct KeysEvent
    keys::Set{AbstractPlotting.Keyboard.Button}
end

@Layoutable Axis begin
    scene::Scene
    xaxislinks::Vector{Axis}
    yaxislinks::Vector{Axis}
    limits::Node{FRect2D}
    block_limit_linking::Node{Bool}
    mouseeventhandle::MouseEventHandle
    scrollevents::Observable{ScrollEvent}
    keysevents::Observable{KeysEvent}
    interactions::Dict{Symbol, Tuple{Bool, Any}}
end

@Layoutable Colorbar

@Layoutable Label

@Layoutable Box

@Layoutable Slider

@Layoutable Button

@Layoutable Toggle

@Layoutable Menu

abstract type LegendElement end

struct LineElement <: LegendElement
    attributes::Attributes
end

struct MarkerElement <: LegendElement
    attributes::Attributes
end

struct PolyElement <: LegendElement
    attributes::Attributes
end

struct LegendEntry
    elements::Vector{LegendElement}
    attributes::Attributes
end

const EntryGroup = Tuple{Optional{String}, Vector{LegendEntry}}

@Layoutable Legend begin
    entrygroups::Node{Vector{EntryGroup}}
end

@Layoutable LScene begin
    scene::Scene
end

@Layoutable Textbox begin
    cursorindex::Node{Int}
    cursoranimtask
end


module MLJFlux

export CUDALibs, CPU1

import Flux
using MLJModelInterface
using MLJModelInterface.ScientificTypesBase
import Base.==
using ProgressMeter
using CategoricalArrays
using Tables
using Statistics
using ColorTypes
using ComputationalResources
using Random

include("penalizers.jl")
include("core.jl")
include("builders.jl")
include("types.jl")
include("regressor.jl")
include("classifier.jl")
include("image.jl")
include("mlj_model_interface.jl")

### Package specific model traits:
MLJModelInterface.metadata_pkg.((NeuralNetworkRegressor,
                                 MultitargetNeuralNetworkRegressor,
                                 NeuralNetworkClassifier,
                                 ImageClassifier),
              name="MLJFlux",
              uuid="094fc8d1-fd35-5302-93ea-dabda2abf845",
              url="https://github.com/alan-turing-institute/MLJFlux.jl",
              julia=true,
              license="MIT")

export NeuralNetworkRegressor, MultitargetNeuralNetworkRegressor
export NeuralNetworkClassifier, ImageClassifier

end #module


#=
This file is auto-generated. Do not edit.
=#
"""
    mutable struct Source <: StaticInjection
        name::String
        available::Bool
        bus::Bus
        V_R::Float64
        V_I::Float64
        X_th::Float64
        services::Vector{Service}
        ext::Dict{String, Any}
        internal::InfrastructureSystemsInternal
    end

This struct acts as an infinity bus.

# Arguments
- `name::String`
- `available::Bool`
- `bus::Bus`
- `V_R::Float64`: Voltage Source Real Component, validation range: (0, nothing)
- `V_I::Float64`: Voltage Source Imaginary Component, validation range: (0, nothing)
- `X_th::Float64`: Source Thevenin impedance, validation range: (0, nothing)
- `services::Vector{Service}`: Services that this device contributes to
- `ext::Dict{String, Any}`
- `internal::InfrastructureSystemsInternal`: power system internal reference, do not modify
"""
mutable struct Source <: StaticInjection
    name::String
    available::Bool
    bus::Bus
    "Voltage Source Real Component"
    V_R::Float64
    "Voltage Source Imaginary Component"
    V_I::Float64
    "Source Thevenin impedance"
    X_th::Float64
    "Services that this device contributes to"
    services::Vector{Service}
    ext::Dict{String, Any}
    "power system internal reference, do not modify"
    internal::InfrastructureSystemsInternal
end

function Source(name, available, bus, V_R, V_I, X_th, services=Device[], ext=Dict{String, Any}(), )
    Source(name, available, bus, V_R, V_I, X_th, services, ext, InfrastructureSystemsInternal(), )
end

function Source(; name, available, bus, V_R, V_I, X_th, services=Device[], ext=Dict{String, Any}(), )
    Source(name, available, bus, V_R, V_I, X_th, services, ext, )
end

# Constructor for demo purposes; non-functional.
function Source(::Nothing)
    Source(;
        name="init",
        available=false,
        bus=Bus(nothing),
        V_R=0,
        V_I=0,
        X_th=0,
        services=Device[],
        ext=Dict{String, Any}(),
    )
end

InfrastructureSystems.get_name(value::Source) = value.name
"""Get Source available."""
get_available(value::Source) = value.available
"""Get Source bus."""
get_bus(value::Source) = value.bus
"""Get Source V_R."""
get_V_R(value::Source) = value.V_R
"""Get Source V_I."""
get_V_I(value::Source) = value.V_I
"""Get Source X_th."""
get_X_th(value::Source) = value.X_th
"""Get Source services."""
get_services(value::Source) = value.services
"""Get Source ext."""
get_ext(value::Source) = value.ext
"""Get Source internal."""
get_internal(value::Source) = value.internal


# 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)...)

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.ϵ)
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.ϵ)
end

function Zygote.accum(x::LennardJones{S, C, W, F, E}, y::LennardJones{S, C, W, F, E}) where {S, C, W, F, E}
    LennardJones{S, C, W, F, E}(x.cutoff, x.nl_only, x.lorentz_mixing, x.weight_14 + y.weight_14, 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, F, E}, CoulombReactionField{D, SO, W, T, F, E}}}; dims=:) where {S, C, W, F, E, D, SO, T}
    reduce(Zygote.accum, xs, dims=dims; init=(
        LennardJones{S, C, W, F, E}(nothing, false, false, zero(W), NoUnits, NoUnits),
        CoulombReactionField{D, SO, W, T, F, E}(zero(D), zero(S), false, zero(W), zero(T), NoUnits, NoUnits),
    ))
end

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

Base.zero(::Type{Atom{T, T, T, T}}) where {T} = Atom(0, zero(T), zero(T), zero(T), zero(T))
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)))
end

function Base.zero(::Type{Union{Nothing, SizedVector{D, T, Vector{T}}}}) where {D, T}
    zero(SizedVector{D, T, Vector{T}})
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 = Zygote._zero(Array(x), eltype(dy))
    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

# Space for 4 duals given to interactions though only one used in this case
# No gradient for cutoff type
function dualize_fb(inter::LennardJones{S, C, W, F, E}) where {S, C, W, F, E}
    w14 = inter.weight_14
    dual_weight_14 = @dualize(w14, 21, 1)
    return LennardJones{S, C, typeof(dual_weight_14), F, E}(inter.cutoff, inter.nl_only,
                inter.lorentz_mixing, dual_weight_14, 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 , 21, 1)
    dual_solvent_dielectric = @dualize(sd , 21, 2)
    dual_weight_14          = @dualize(w14, 21, 3)
    dual_coulomb_const      = @dualize(cc , 21, 4)
    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)
end

function dualize_atom_fb1(at::Atom)
    charge, mass, σ, ϵ = at.charge, at.mass, at.σ, at.ϵ
    dual_charge = @dualize(charge, 21, 11)
    dual_mass = @dualize(mass, 21, 12)
    dual_σ = @dualize(σ, 21, 13)
    dual_ϵ = @dualize(ϵ, 21, 14)
    return Atom{typeof(dual_charge), typeof(dual_mass), typeof(dual_σ), typeof(dual_ϵ)}(
                at.index, dual_charge, dual_mass, dual_σ, dual_ϵ)
end

function dualize_atom_fb2(at::Atom)
    charge, mass, σ, ϵ = at.charge, at.mass, at.σ, at.ϵ
    dual_charge = @dualize(charge, 21, 15)
    dual_mass = @dualize(mass, 21, 16)
    dual_σ = @dualize(σ, 21, 17)
    dual_ϵ = @dualize(ϵ, 21, 18)
    return Atom{typeof(dual_charge), typeof(dual_mass), typeof(dual_σ), typeof(dual_ϵ)}(
                at.index, dual_charge, dual_mass, dual_σ, dual_ϵ)
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(isa(arg2, 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)
        charge, mass, σ, ϵ = arg1.charge, arg1.mass, arg1.σ, arg1.ϵ
        ds1 = Atom(arg1.index, @dualize(charge, 4, 1), @dualize(mass, 4, 2),
                    @dualize(σ, 4, 3), @dualize(ϵ, 4, 4))
        return f(ds1)
    end
end

function dual_function_force_broadcast(f::F) where F
    function (arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8)
        ds1 = (dualize_fb(arg1[1]), dualize_fb(arg1[2])) # Only works for this case
        ds2 = dualize(Nothing, arg2, Val(4), Val(14))
        ds3 = dualize(Nothing, arg3, Val(7), Val(11))
        ds4 = dualize_atom_fb1(arg4)
        ds5 = dualize_atom_fb2(arg5)
        ds6 = dualize(Nothing, arg6, Val(18), 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

@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

@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)
        ȳ = arg1 isa CuArray ? cu(ȳ_in) : ȳ_in
        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)
        ȳ = arg1 isa CuArray ? cu(ȳ_in) : ȳ_in
        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)
        ȳ = arg1 isa CuArray ? cu(ȳ_in) : ȳ_in
        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

@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)
        ȳ = arg1 isa CuArray ? cu(ȳ_in) : ȳ_in
        barg1 = broadcast(ȳ, out) do y1, o1
            ps = partials(o1)
            Atom(0, y1 * ps[1], y1 * ps[2], y1 * ps[3], y1 * ps[4])
        end
        darg1 = unbroadcast(arg1, barg1)
        (nothing, nothing, darg1)
    end
    return y, bc_fwd_back
end

function combine_dual_GeneralInteraction(y1::SVector{3, T}, o1::SVector{3, Dual{Nothing, T, P}}, i::Integer) where {T, P}
    (
        LennardJones{false, Nothing, 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),
            NoUnits, NoUnits,
        ),
        CoulombReactionField{T, T, T, T, typeof(NoUnits), typeof(NoUnits)}(
            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),
            false,
            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),
            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),
        ),
    )
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],
    )
end

@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)
        ȳ = arg2 isa CuArray ? cu(ȳ_in) : ȳ_in
        darg1 = unbroadcast(arg1, broadcast(combine_dual_GeneralInteraction, ȳ, out, 1))
        darg2 = unbroadcast(arg2, broadcast((y1, o1) -> SVector{D, T}(sum_partials(o1, y1,  5),
                                sum_partials(o1, y1,  6), sum_partials(o1, y1,  7)), ȳ, out))
        darg3 = unbroadcast(arg3, broadcast((y1, o1) -> SVector{D, T}(sum_partials(o1, y1,  8),
                                sum_partials(o1, y1,  9), sum_partials(o1, y1, 10)), ȳ, out))
        darg4 = unbroadcast(arg4, broadcast(combine_dual_Atom, ȳ, out, 11, 12, 13, 14))
        darg5 = unbroadcast(arg5, broadcast(combine_dual_Atom, ȳ, out, 15, 16, 17, 18))
        darg6 = unbroadcast(arg6, broadcast((y1, o1) -> SVector{D, T}(sum_partials(o1, y1, 19),
                                sum_partials(o1, y1, 20), sum_partials(o1, y1, 21)), ȳ, out))
        darg7 = nothing
        darg8 = nothing
        return (nothing, nothing, darg1, darg2, darg3, darg4, darg5, darg6, darg7, darg8)
    end
    return y, bc_fwd_back
end

@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)
        ȳ = arg1 isa CuArray ? cu(ȳ_in) : ȳ_in
        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

@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)
        ȳ = arg1 isa CuArray ? cu(ȳ_in) : ȳ_in
        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

@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)
        ȳ = arg1 isa CuArray ? cu(ȳ_in) : ȳ_in
        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

@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 (:+, :-, :*, :/, :force, :force_nounit, :mass, :remove_molar, :ustrip,
            :ustrip_vec, :wrap_coords_vec, :getf1, :getf2, :getf3, :getf4)
    @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